vnstat module

1 files changed, 138 insertions, 0 deletions

diff --git a/modules/vnstat.py b/modules/vnstat.py
new file mode 100644
index 0000000..100138c
--- /dev/null
+++ b/modules/vnstat.py
@@ -0,0 +1,138 @@
+# -*- coding: utf8 -*-
+
+from __future__ import division  # python2 compatibility
+from time import time, sleep
+from subprocess import check_output
+
+
+"""
+Module for displaying vnstat's statistics.
+REQUIRE external program called "vnstat" installed and configured to work.
+
+@author shadowprince
+@license Eclipse Public License
+"""
+
+CACHED_TIME = 3*60  # update time (in seconds)
+POSITION = 0  # bar position
+STATISTICS_TYPE = "d"  # d for daily, m for monthly
+
+"""
+Coloring rules.
+
+If value is bigger that dict key, status string will turn to color, specified in the value.
+Example:
+COLORING = {
+    800: "#dddd00",
+    900: "#dd0000",
+}
+(0 - 800: white, 800-900: yellow, >900 - red)
+"""
+COLORING = {}
+
+"""
+Format of status string.
+
+Placeholders:
+    total - total
+    up - upload
+    down - download
+"""
+FORMAT = "{total}"
+
+PRECISION = 1  # amount of numbers after dot
+MULTIPLIER_TOP = 1024  # if value is greater, divide it with UNIT_MULTI and get next unit from UNITS
+LEFT_ALIGN = 0
+
+"""
+Format of total, up and down placeholders under FORMAT.
+As default, substitutes LEFT_ALIGN and PRECISION as %s and %s
+Placeholders:
+    value - value (float)
+    unit - unit (string)
+
+"""
+VALUE_FORMAT = "{value:%s.%sf} {unit}" % (LEFT_ALIGN, PRECISION)
+
+INITIAL_MULTI = 1024  # initial multiplier, if you want to get rid of first bytes, set to 1 to disable
+UNIT_MULTI = 1024  # value to divide if rate is greater than MULTIPLIER_TOP
+UNITS = ["kb", "mb", "gb", "tb", ]  # list of units, first one - value/INITIAL_MULTI, second - value/1024, third - value/1024^2, etc...
+
+
+def get_stat():
+    """
+    Get statistics from devfile in list of lists of words
+    """
+    def filter_stat():
+        for x in check_output(["vnstat", "--dumpdb"]).decode("utf-8").splitlines():
+            if x.startswith("{};0;".format(STATISTICS_TYPE)):
+                return x
+
+    try:
+        type, number, ts, rxm, txm, rxk, txk, fill = filter_stat().split(";")
+    except OSError as e:
+        print("Looks like you have'nt installed or configured vnstat!")
+        raise e
+    except ValueError:
+        raise RuntimeError("vnstat returned wrong output, maybe it's configured wrong or module is outdated")
+
+    up = (int(txm) * 1024 + int(txk)) * 1024
+    down = (int(rxm) * 1024 + int(rxk)) * 1024
+
+    return {"up": up,
+            "down": down,
+            "total": up+down, }
+
+
+def divide_and_format(value):
+    """
+    Divide a value and return formatted string
+    """
+    value /= INITIAL_MULTI
+    for i, unit in enumerate(UNITS):
+        if value > MULTIPLIER_TOP:
+            value /= UNIT_MULTI
+        else:
+            break
+
+    return VALUE_FORMAT.format(value=value, unit=unit)
+
+
+class Py3status:
+    def __init__(self, *args, **kwargs):
+        self.last_stat = get_stat()
+        self.last_time = time()
+        self.last_interface = None
+
+    def currentSpeed(self, json, i3status_config):
+        stat = get_stat()
+
+        color = None
+        keys = list(COLORING.keys())
+        keys.sort()
+        for k in keys:
+            if stat["total"] < k * 1024 * 1024:
+                break
+            else:
+                color = COLORING[k]
+
+        response = {
+            'transformed': True,
+            'full_text': FORMAT.format(
+                total=divide_and_format(stat['total']),
+                up=divide_and_format(stat['up']),
+                down=divide_and_format(stat['down']),),
+            'name': 'vnstat',
+            'cached_until': time() + CACHED_TIME,
+        }
+
+        if color:
+            response["color"] = color
+
+        return POSITION, response
+
+if __name__ == "__main__":
+    x = Py3status()
+    while True:
+        print(x.currentSpeed(1, 1))
+        sleep(1)