# pylint: disable=C0111,R0903

"""Displays network IO for interfaces.

Parameters:
    * traffic.exclude: Comma-separated list of interface prefixes to exclude (defaults to "lo,virbr,docker,vboxnet,veth")
"""

import re
import psutil
import netifaces

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

class Module(bumblebee.engine.Module):
    def __init__(self, engine, config):
        widgets = []
        super(Module, self).__init__(engine, config, widgets)
        self._exclude = tuple(filter(len, self.parameter("exclude", "lo,virbr,docker,vboxnet,veth").split(",")))
        self._update_widgets(widgets)
        self._status = ""

    def state(self, widget):
        if "traffic.rx" in widget.name:
            return "rx"
        if "traffic.tx" in widget.name:
            return "tx"
        return self._status

    def update(self, widgets):
        self._update_widgets(widgets)

    def create_widget(self, widgets, name, txt=None, attributes={}):
        widget = self.widget(name)
        if widget: return widget

        widget = bumblebee.output.Widget(name=name)
        widget.full_text(txt)
        widgets.append(widget)

        for key in attributes:
            widget.set(key, attributes[key])

        return widget

    def _update_widgets(self, widgets):
        interfaces = [ i for i in netifaces.interfaces() if not i.startswith(self._exclude) ]

        counters = psutil.net_io_counters(pernic=True)
        for interface in interfaces:
            if not interface: interface = "lo"
            data = {
                "rx": counters[interface].bytes_recv,
                "tx": counters[interface].bytes_sent,
            }

            name = "traffic-{}".format(interface)

            self.create_widget(widgets, name, interface)

            for direction in ["rx", "tx"]:
                name = "traffic.{}-{}".format(direction, interface)
                widget = self.create_widget(widgets, name, attributes={"theme.minwidth": "1000.00MB"})
                prev = widget.get(direction, 0)
                speed = bumblebee.util.bytefmt(int(data[direction]) - int(prev))
                widget.full_text(speed)
                widget.set(direction, data[direction])

# vim: tabstop=8 expandtab shiftwidth=4 softtabstop=4

@@ -198,16 +198,165 @@
 ,veth%22)%0A
+    * traffic.states: Comma-separated list of states to show (prefix with %22%5E%22 to invert - i.e. %5Edown -%3E show all devices that are not in state down)%0A
 %22%22%22%0A%0Aimp
@@ -763,16 +763,387 @@
 (%22,%22)))%0A
+        self._status = %22%22%0A%0A        self._states = %7B%7D%0A        self._states%5B%22include%22%5D = %5B%5D%0A        self._states%5B%22exclude%22%5D = %5B%5D%0A        for state in tuple(filter(len, self.parameter(%22states%22, %22%22).split(%22,%22))):%0A            if state%5B0%5D == %22%5E%22:%0A                self._states%5B%22exclude%22%5D.append(state%5B1:%5D)%0A            else:%0A                self._states%5B%22include%22%5D.append(state)%0A
         
@@ -1175,34 +1175,8 @@
 ets)
-%0A        self._status = %22%22
 %0A%0A  
@@ -1793,273 +1793,892 @@
 def 
-_update_widgets(self, widgets):%0A        interfaces = %5B i for i in netifaces.interfaces() if not i.startswith(self._exclude) %5D%0A%0A        counters = psutil.net_io_counters(pernic=True)%0A        for interface in interfaces:%0A            if not interface: interface = %22lo%22
+get_addresses(self, intf):%0A        retval = %5B%5D%0A        try:%0A            for ip in netifaces.ifaddresses(intf).get(netifaces.AF_INET, %5B%5D):%0A                if ip.get(%22addr%22, %22%22) != %22%22:%0A                    retval.append(ip.get(%22addr%22))%0A        except Exception:%0A            return %5B%5D%0A        return retval%0A%0A    def _update_widgets(self, widgets):%0A        interfaces = %5B i for i in netifaces.interfaces() if not i.startswith(self._exclude) %5D%0A%0A        counters = psutil.net_io_counters(pernic=True)%0A        for interface in interfaces:%0A            if not interface: interface = %22lo%22%0A            state = %22down%22%0A            if len(self.get_addresses(interface)) %3E 0:%0A                state = %22up%22%0A%0A            if len(self._states%5B%22exclude%22%5D) %3E 0 and state in self._states%5B%22exclude%22%5D: continue%0A            if len(self._states%5B%22include%22%5D) %3E 0 and state not in self._states%5B%22include%22%5D: continue%0A
 %0A   

Allow for state filtering