bumblebee-status/bumblebee/modules/network_traffic.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""Displays network traffic
   * No extra configuration needed
"""

import psutil
import netifaces

import bumblebee.input
import bumblebee.output
import bumblebee.engine
import bumblebee.util

WIDGET_NAME = 'network_traffic'

class Module(bumblebee.engine.Module):
    """Bumblebee main module """

    def __init__(self, engine, config):
        super(Module, self).__init__(engine, config)

        try:
            self._bandwidth = BandwidthInfo()

            self._download_tx = self._bandwidth.bytes_recv()
            self._upload_tx = self._bandwidth.bytes_sent()
        except Exception:
            """ We do not want do explode anything """
            pass

    @classmethod
    def state(cls, widget):
        """Return the widget state"""

        if widget.name == '{}.rx'.format(WIDGET_NAME):
            return 'rx'
        elif widget.name == '{}.tx'.format(WIDGET_NAME):
            return 'tx'

        return None

    def update(self, widgets):
        try:
            download_tx = self._bandwidth.bytes_recv()
            upload_tx = self._bandwidth.bytes_sent()

            download_rate = (download_tx - self._download_tx)
            upload_rate = (upload_tx - self._upload_tx)

            self.update_widgets(widgets, download_rate, upload_rate)

            self._download_tx, self._upload_tx = download_tx, upload_tx
        except Exception:
            """ We do not want do explode anything """
            pass

    @classmethod
    def update_widgets(cls, widgets, download_rate, upload_rate):
        """Update tx/rx widgets with new rates"""
        del widgets[:]

        widgets.extend((
            TrafficWidget(text=download_rate, direction='tx'),
            TrafficWidget(text=upload_rate, direction='rx')
        ))


class BandwidthInfo(object):
    """Get received/sent bytes from network adapter"""

    def bytes_recv(self):
        """Return received bytes"""
        return self.bandwidth().bytes_recv

    def bytes_sent(self):
        """Return sent bytes"""
        return self.bandwidth().bytes_sent

    def bandwidth(self):
        """Return bandwidth information"""
        io_counters = self.io_counters()
        return io_counters[self.default_network_adapter()]

    @classmethod
    def default_network_adapter(cls):
        """Return default active network adapter"""
        gateway = netifaces.gateways()['default']

        if not gateway:
            raise 'No default gateway found'

        return gateway[netifaces.AF_INET][1]

    @classmethod
    def io_counters(cls):
        """Return IO counters"""
        return psutil.net_io_counters(pernic=True)

class TrafficWidget(object):
    """Create a traffic widget with humanized bytes string with proper icon (up/down)"""
    def __new__(cls, text, direction):
        widget = bumblebee.output.Widget(name='{0}.{1}'.format(WIDGET_NAME, direction))
        widget.set('theme.minwidth', '0000000KiB/s')
        widget.full_text(cls.humanize(text))

        return widget

    @staticmethod
    def humanize(text):
        """Return humanized bytes"""
        humanized_byte_format = bumblebee.util.bytefmt(text)
        return '{0}/s'.format(humanized_byte_format)