#!/usr/bin/env python

# Copyright (C) 2013 Citrix Systems Inc.
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as published
# by the Free Software Foundation; version 2.1 only. with the special
# exception on linking described in file LICENSE.
#
# This program 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 Lesser General Public License for more details.

# xsifstat : Report per-VIF network throughput figures, and aggregate
#            values for each netback thread and bridge device.
#            Note that 'Tx' and 'Rx' are from the perspective of dom0:
#            - Rx is bytes received _by_ dom0, i.e. transmitted by guest
#            - Tx is bytes transmitted _by_ dom0, i.e. received by guest
#            This method of reporting is used by ifconfig, for instance,
#            so we use it here for consistency across tools.

import re
import signal
import sys
import time

sample_period = 2 # Seconds

# Register a handler for SIGINT to exit cleanly without a stack trace
def signal_handler(sig, frame):
    sys.exit(0)
signal.signal(signal.SIGINT, signal_handler)

# Build a regular expression to extract information from
# /proc/net/dev
# Only need to capture the VIF name, rx bytes and tx bytes
numeric_field = r'\s*\d+'
numeric_field_cap = r'\s*(\d+)'
vif_label_cap = r'\s*(vif\d+.\d+):'
regex = re.compile(
        vif_label_cap # Interface name
        + numeric_field_cap # Rx bytes
        + numeric_field # Rx packets
        + numeric_field # Rx errs
        + numeric_field # Rx drop
        + numeric_field # Rx fifo
        + numeric_field # Rx frame
        + numeric_field # Rx compressed
        + numeric_field # Rx multicast
        + numeric_field_cap # Tx bytes
        + numeric_field # Tx packets
        + numeric_field # Tx errs
        + numeric_field # Tx drop
        + numeric_field # Tx fifo
        + numeric_field # Tx colls
        + numeric_field # Tx carrier
        + numeric_field # Tx compressed
        )

cache = { } # Holds mapping vif -> (rxbytes, txbytes, bridge) tuple

current = { } # Holds information about the current sample
# Structure:
# {bridge_name -> {
#   'tx' -> int,
#   'rx' -> int,
#   'children' -> {
#     vif_name -> (dtx, drx)}}}}}

# Do the open once, outside the loop and seek back to the beginning
# for every read (thus avoiding closing & reopening)
try:
    devfile = open('/proc/net/dev', 'r')
except:
    print 'Unable to open /proc/net/dev. Exiting.'
    sys.exit(1)

attached_vifs = set() # unattached vifs are periodically removed from cache

while True:
    devfile.seek(0) # Read the file from the beginning
    current = { } # Reset current data map
    # Remove expired VIFs from cache
    unattached_vifs = set(cache.keys()) - attached_vifs
    for vif in unattached_vifs:
        del cache[vif]
    attached_vifs = set()
    # Process the information about attached VIFs
    for line in devfile:
        match = regex.search(line)
        if match:
            vif = match.group(1)
            rx_bytes = int(match.group(2))
            tx_bytes = int(match.group(3))
            bridge = ''
            if vif in cache:
                # Have existing information for this interface, so we
                # can look up the bridge, find the
                # previous Tx and Rx bytes, and report a throughput
                # this time around
                bridge = cache[vif][2]
                drx = rx_bytes - cache[vif][0]
                # Handle the case where the tx or rx bytes wrap around
                if drx < 0:
                    drx = 2**32 - cache[vif][0] + rx_bytes 
                dtx = tx_bytes - cache[vif][1]
                if dtx < 0:
                    dtx = 2**32 - cache[vif][1] + tx_bytes
                # Add this VIF to the current report
                if bridge not in current:
                    current[bridge] = { 'tx': 0.0, 'rx': 0.0, 'children': {} }
                current[bridge]['tx'] += dtx
                current[bridge]['rx'] += drx
                current[bridge]['children'][vif] = (drx, dtx)
            else:
                try:
                    # Determine bridge
                    f = open('/sys/class/net/%s/brport/bridge/uevent'
                            % vif, 'r')
                    iflines = [line for line in f.readlines() if line.startswith("INTERFACE=")]
                    bridge = iflines[0].replace('INTERFACE=','').strip()
                    f.close()
                    attached_vifs.add(vif) # Add this VIF to the attached list
                    # This is done so that we don't immediately remove it
                    # from the cache, which would mean it never got displayed
                except:
                    pass # Ignore the case where we can't read the file above
            # Store all the information in the cache so it is available next
            # sample period
            cache[vif] = (rx_bytes, tx_bytes, bridge)
    # Output any stats we have
    for bridge, bridgestats in current.iteritems():
        bridge_rxrate = float(bridgestats['rx']*8) / float(sample_period)
        bridge_rxrate /= 1e9 # Gbit/s
        bridge_txrate = float(bridgestats['tx']*8) / float(sample_period)
        bridge_txrate /= 1e9 # Gbit/s
        print '\n%s: Rx=%.3f Gbit/s, Tx=%.3f Gbit/s' \
                % (bridge, bridge_rxrate, bridge_txrate)
        for vif, vifstats in bridgestats['children'].iteritems():
            attached_vifs.add(vif) # Add this vif to the list of attached vifs
            rxrate = float(vifstats[0]*8) / float(sample_period)
            rxrate /= 1e9 # Gbit/s
            txrate = float(vifstats[1]*8) / float(sample_period)
            txrate /= 1e9 # Gbit/s
            print "|- %s: Rx=%.3f Gbit/s, Tx=%.3f Gbit/s" \
                    % (vif, rxrate, txrate)
    # Now wait sample_period seconds
    time.sleep(sample_period)

