[feature] add up/down traffic in status

7 files changed, 38 insertions, 20 deletions

diff --git a/src/leap/bitmask/vpn/_control.py b/src/leap/bitmask/vpn/_control.py
index e05aeed..31a1216 100644
--- a/src/leap/bitmask/vpn/_control.py
+++ b/src/leap/bitmask/vpn/_control.py
@@ -98,6 +98,10 @@ class VPNControl(object):
             return 'OFFLINE'
         return self._vpnproc.status
 
+    @property
+    def traffic_status(self):
+        return self._vpnproc.traffic_status
+
 
     def stop(self, shutdown=False, restart=False):
         """
diff --git a/src/leap/bitmask/vpn/_human.py b/src/leap/bitmask/vpn/_human.py
index 1b2ba64..803d365 100644
--- a/src/leap/bitmask/vpn/_human.py
+++ b/src/leap/bitmask/vpn/_human.py
@@ -1,6 +1,3 @@
-## {{{ http://code.activestate.com/recipes/578019/ (r15)
-#!/usr/bin/env python
-
 """
 Bytes-to-human / human-to-bytes converter.
 Based on: http://goo.gl/kTQMs
@@ -12,14 +9,15 @@ License: MIT
 
 # see: http://goo.gl/kTQMs
 SYMBOLS = {
-    'customary'     : ('B', 'K', 'M', 'G', 'T', 'P', 'E', 'Z', 'Y'),
-    'customary_ext' : ('byte', 'kilo', 'mega', 'giga', 'tera', 'peta', 'exa',
-                       'zetta', 'iotta'),
-    'iec'           : ('Bi', 'Ki', 'Mi', 'Gi', 'Ti', 'Pi', 'Ei', 'Zi', 'Yi'),
-    'iec_ext'       : ('byte', 'kibi', 'mebi', 'gibi', 'tebi', 'pebi', 'exbi',
-                       'zebi', 'yobi'),
+    'customary': ('B', 'K', 'M', 'G', 'T', 'P', 'E', 'Z', 'Y'),
+    'customary_ext': ('byte', 'kilo', 'mega', 'giga', 'tera', 'peta', 'exa',
+                      'zetta', 'iotta'),
+    'iec': ('Bi', 'Ki', 'Mi', 'Gi', 'Ti', 'Pi', 'Ei', 'Zi', 'Yi'),
+    'iec_ext': ('byte', 'kibi', 'mebi', 'gibi', 'tebi', 'pebi', 'exbi',
+                'zebi', 'yobi'),
 }
 
+
 def bytes2human(n, format='%(value).1f %(symbol)s', symbols='customary'):
     """
     Convert n bytes into a human readable string based on format.
@@ -63,13 +61,14 @@ def bytes2human(n, format='%(value).1f %(symbol)s', symbols='customary'):
     symbols = SYMBOLS[symbols]
     prefix = {}
     for i, s in enumerate(symbols[1:]):
-        prefix[s] = 1 << (i+1)*10
+        prefix[s] = 1 << (i + 1) * 10
     for symbol in reversed(symbols[1:]):
         if n >= prefix[symbol]:
             value = float(n) / prefix[symbol]
             return format % locals()
     return format % dict(symbol=symbols[0], value=n)
 
+
 def human2bytes(s):
     """
     Attempts to guess the string format based on default symbols
@@ -115,8 +114,7 @@ def human2bytes(s):
             letter = letter.upper()
         else:
             raise ValueError("can't interpret %r" % init)
-    prefix = {sset[0]:1}
+    prefix = {sset[0]: 1}
     for i, s in enumerate(sset[1:]):
-        prefix[s] = 1 << (i+1)*10
+        prefix[s] = 1 << (i + 1) * 10
     return int(num * prefix[letter])
-
diff --git a/src/leap/bitmask/vpn/_management.py b/src/leap/bitmask/vpn/_management.py
index 905372b..8bf3cc7 100644
--- a/src/leap/bitmask/vpn/_management.py
+++ b/src/leap/bitmask/vpn/_management.py
@@ -265,7 +265,11 @@ class VPNManagement(object):
             if len(parts) < 2:
                 continue
 
-            text, value = parts
+            try:
+                text, value = parts
+            except ValueError:
+                logger.debug("Could not parse %s" % parts)
+                return
             # text can be:
             #   "TUN/TAP read bytes"
             #   "TUN/TAP write bytes"
diff --git a/src/leap/bitmask/vpn/_status.py b/src/leap/bitmask/vpn/_status.py
index 11d564f..da8992e 100644
--- a/src/leap/bitmask/vpn/_status.py
+++ b/src/leap/bitmask/vpn/_status.py
@@ -1,9 +1,11 @@
 from itertools import chain, repeat
 from twisted.logger import Logger
+from ._human import bytes2human
 
 logger = Logger()
 
 
+
 # TODO implement a state machine in this class
 
 
@@ -49,20 +51,18 @@ class VPNStatus(object):
         status, errcode = self._status_codes(event)
         self.set_status(status, errcode)
 
-
     def set_status(self, status, errcode):
         self.status = status
         self.errcode = errcode
 
     def set_traffic_status(self, status):
-        down, up = status
+        up, down = status
         self._traffic_up = up
         self._traffic_down = down
 
     def get_traffic_status(self):
-        # XXX return Human readable too
-        return {'down': self._traffic_down,
-                'up': self._traffic_up}
+        return {'down': bytes2human(self._traffic_down),
+                'up': bytes2human(self._traffic_up)}
 
     def _status_codes(self, event):
         # TODO check good transitions
diff --git a/src/leap/bitmask/vpn/eip.py b/src/leap/bitmask/vpn/eip.py
index c2aa4fb..0c47cc3 100644
--- a/src/leap/bitmask/vpn/eip.py
+++ b/src/leap/bitmask/vpn/eip.py
@@ -75,7 +75,11 @@ class EIPManager(object):
     def get_status(self):
         vpn_status = self._vpn.status
         fw_status = self._firewall.status
+
         result = {'EIP': vpn_status,
                   'firewall': fw_status}
+        if vpn_status == 'CONNECTED':
+            traffic = self._vpn.traffic_status
+            result['↑↑↑'] = traffic['up']
+            result['↓↓↓'] = traffic['down']
         return result
-
diff --git a/src/leap/bitmask/vpn/manager.py b/src/leap/bitmask/vpn/manager.py
index 04e8809..8781603 100644
--- a/src/leap/bitmask/vpn/manager.py
+++ b/src/leap/bitmask/vpn/manager.py
@@ -89,6 +89,10 @@ class VPNManager(object):
     def status(self):
         return self._vpn.status
 
+    @property
+    def traffic_status(self):
+        return self._vpn.traffic_status
+
 
     def _get_management_location(self):
         """
diff --git a/src/leap/bitmask/vpn/process.py b/src/leap/bitmask/vpn/process.py
index 185ba62..82b4e7c 100644
--- a/src/leap/bitmask/vpn/process.py
+++ b/src/leap/bitmask/vpn/process.py
@@ -103,6 +103,10 @@ class VPNProcess(protocol.ProcessProtocol, _management.VPNManagement):
     def status(self):
         return self._status.status
 
+    @property
+    def traffic_status(self):
+        return self._status.get_traffic_status()
+
     # processProtocol methods
 
     def connectionMade(self):