[modules] Re-enable traffic module

modules/contrib/traffic.py

@@ -0,0 +1,159 @@
+# 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")
+    * traffic.states: Comma-separated list of states to show (prefix with "^" to invert - i.e. ^down -> show all devices that are not in state down)
+    * traffic.showname: If set to False, hide network interface name (defaults to True)
+    * traffic.format: Format string for download/upload speeds.
+                      Defaults to "{:.2f}"
+    * traffic.graphlen: Graph lenth in seconds. Positive even integer. Each
+                        char shows 2 seconds. If set, enables up/down traffic
+                        graphs
+"""
+
+import re
+import time
+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._status = ""
+
+        self._showname = bumblebee.util.asbool(self.parameter("showname", True))
+        self._format = self.parameter("format", "{:.2f}")
+        self._prev = {}
+        self._states = {}
+        self._lastcheck = 0
+        self._states["include"] = []
+        self._states["exclude"] = []
+        for state in tuple(filter(len, self.parameter("states", "").split(","))):
+            if state[0] == "^":
+                self._states["exclude"].append(state[1:])
+            else:
+                self._states["include"].append(state)
+        self._graphlen = int(self.parameter("graphlen", 0))
+        if self._graphlen > 0:
+            self._graphdata = {}
+        self._first_run = True
+        self._update_widgets(widgets)
+
+    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 = bumblebee.output.Widget(name=name)
+        widget.full_text(txt)
+        widgets.append(widget)
+
+        for key in attributes:
+            widget.set(key, attributes[key])
+
+        return widget
+
+    def get_addresses(self, intf):
+        retval = []
+        try:
+            for ip in netifaces.ifaddresses(intf).get(netifaces.AF_INET, []):
+                if ip.get("addr", "") != "":
+                    retval.append(ip.get("addr"))
+        except Exception:
+            return []
+        return retval
+
+    def get_minwidth_str(self):
+        """
+            computes theme.minwidth string
+            based on traffic.format and traffic.graphlen parameters
+        """
+        minwidth_str = ""
+        if self._graphlen > 0:
+            graph_len = int(self._graphlen / 2)
+            graph_prefix = "0" * graph_len
+            minwidth_str += graph_prefix
+        minwidth_str += "1000"
+        try:
+            length = int(re.match("{:\.(\d+)f}", self._format).group(1))
+            if length > 0:
+                minwidth_str += "." + "0" * length
+        except AttributeError:
+            # return default value
+            return "1000.00KiB/s"
+        finally:
+            minwidth_str += "KiB/s"
+        return minwidth_str
+
+    def _update_widgets(self, widgets):
+        interfaces = [i for i in netifaces.interfaces() if not i.startswith(self._exclude)]
+
+        del widgets[:]
+
+        counters = psutil.net_io_counters(pernic=True)
+        now = time.time()
+        timediff = now - (self._lastcheck if self._lastcheck else now)
+        if timediff <= 0: timediff = 1
+        self._lastcheck = now
+        for interface in interfaces:
+            if self._graphlen > 0:
+                if interface not in self._graphdata:
+                    self._graphdata[interface] = {
+                        "rx": [0] * self._graphlen,
+                        "tx": [0] * self._graphlen}
+            if not interface: interface = "lo"
+            state = "down"
+            if len(self.get_addresses(interface)) > 0:
+                state = "up"
+            elif bumblebee.util.asbool(self.parameter("hide_down", True)):
+                continue
+
+            if len(self._states["exclude"]) > 0 and state in self._states["exclude"]: continue
+            if len(self._states["include"]) > 0 and state not in self._states["include"]: continue
+
+            data = {
+                "rx": counters[interface].bytes_recv,
+                "tx": counters[interface].bytes_sent,
+            }
+
+            name = "traffic-{}".format(interface)
+
+            if self._showname:
+               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": self.get_minwidth_str()})
+                prev = self._prev.get(name, 0)
+                bspeed = (int(data[direction]) - int(prev))/timediff
+                speed = bumblebee.util.bytefmt(bspeed, self._format)
+                txtspeed = '{0}/s'.format(speed)
+                if self._graphlen > 0:
+                    # skip first value returned by psutil, because it is
+                    # giant and ruins the grapth ratio until it gets pushed
+                    # out of saved list
+                    if self._first_run is True:
+                        self._first_run = False
+                    else:
+                        self._graphdata[interface][direction] = self._graphdata[interface][direction][1:]
+                        self._graphdata[interface][direction].append(bspeed)
+                    txtspeed = "{}{}".format(bumblebee.output.bgraph(self._graphdata[interface][direction]), txtspeed)
+                widget.full_text(txtspeed) 
+                self._prev[name] = data[direction]
+
+# vim: tabstop=8 expandtab shiftwidth=4 softtabstop=4