[feature] add twisted protocol for handling openvpn management

6 files changed, 644 insertions, 399 deletions

diff --git a/src/leap/bitmask/vpn/_management.py b/src/leap/bitmask/vpn/_management.py
index fac1b09..d05790c 100644
--- a/src/leap/bitmask/vpn/_management.py
+++ b/src/leap/bitmask/vpn/_management.py
@@ -15,7 +15,6 @@ except ImportError:
     from psutil import AccessDenied as psutil_AccessDenied
     PSUTIL_2 = True
 
-from leap.bitmask.vpn._telnet import UDSTelnet
 
 
 class OpenVPNAlreadyRunning(Exception):
@@ -32,237 +31,8 @@ class ImproperlyConfigured(Exception):
     pass
 
 
-class VPNManagement(object):
-    """
-    A class to handle the communication with the openvpn management
-    interface.
-
-    For more info about management methods::
-
-      zcat `dpkg -L openvpn | grep management`
-    """
-    log = Logger()
-
-    # Timers, in secs
-    CONNECTION_RETRY_TIME = 1
-
-    def __init__(self):
-        self._tn = None
-        self.aborted = False
-        self._host = None
-        self._port = None
-
-        self._watcher = None
-        self._logs = {}
-
-    def set_connection(self, host, port):
-        """
-        :param host: either socket path (unix) or socket IP
-        :type host: str
+class Management(object):
 
-        :param port: either string "unix" if it's a unix socket, or port
-                     otherwise
-        """
-        self._host = host
-        self._port = port
-
-    def set_watcher(self, watcher):
-        self._watcher = watcher
-
-    def is_connected(self):
-        return bool(self._tn)
-
-    def connect_retry(self, retry=0, max_retries=None):
-        """
-        Attempts to connect to a management interface, and retries
-        after CONNECTION_RETRY_TIME if not successful.
-
-        :param retry: number of the retry
-        :type retry: int
-        """
-        if max_retries and retry > max_retries:
-            self.log.warn(
-                'Max retries reached while attempting to connect '
-                'to management. Aborting.')
-            self.aborted = True
-            return
-
-        if not self.aborted and not self.is_connected():
-            self._connect()
-            reactor.callLater(
-                self.CONNECTION_RETRY_TIME,
-                self.connect_retry, retry + 1, max_retries)
-
-    def _connect(self):
-        if not self._host or not self._port:
-            raise ImproperlyConfigured('Connection is not configured')
-
-        try:
-            self._tn = UDSTelnet(self._host, self._port)
-            self._tn.read_eager()
-
-        except Exception as e:
-            self.log.warn('Could not connect to OpenVPN yet: %r' % (e,))
-            self._tn = None
-
-        if self._tn:
-            return True
-        else:
-            self.log.error('Error while connecting to management!')
-            return False
-
-    def process_log(self):
-        if not self._watcher or not self._tn:
-            return
-
-        lines = self._send_command('log 20')
-        for line in lines:
-            try:
-                splitted = line.split(',')
-                ts = splitted[0]
-                msg = ','.join(splitted[2:])
-                if msg.startswith('MANAGEMENT'):
-                    continue
-                if ts not in self._logs:
-                    self._watcher.watch(msg)
-                    self.log.info('VPN: %s' % msg)
-                    self._logs[ts] = msg
-            except Exception:
-                pass
-
-    def _seek_to_eof(self):
-        """
-        Read as much as available. Position seek pointer to end of stream
-        """
-        try:
-            return self._tn.read_eager()
-        except EOFError:
-            self.log.debug('Could not read from socket. Assuming it died.')
-            return
-
-    def _send_command(self, command, until=b"END"):
-        """
-        Sends a command to the telnet connection and reads until END
-        is reached.
-
-        :param command: command to send
-        :type command: str
-
-        :param until: byte delimiter string for reading command output
-        :type until: byte str
-
-        :return: response read
-        :rtype: list
-        """
-        try:
-            self._tn.write("%s\n" % (command,))
-            buf = self._tn.read_until(until)
-            seek = self._seek_to_eof()
-            blist = buf.split('\r\n')
-            if blist[-1].startswith(until):
-                del blist[-1]
-                return blist
-            else:
-                return []
-
-        except socket.error:
-            # XXX should get a counter and repeat only
-            # after mod X times.
-            self.log.warn('Socket error (command was: "%s")' % (command,))
-            self._close_management_socket(announce=False)
-            self.log.debug('Trying to connect to management again')
-            self.connect_retry(max_retries=5)
-            return []
-
-        except Exception as e:
-            self.log.warn("Error sending command %s: %r" %
-                          (command, e))
-        return []
-
-    def _close_management_socket(self, announce=True):
-        """
-        Close connection to openvpn management interface.
-        """
-        if announce:
-            self._tn.write("quit\n")
-            self._tn.read_all()
-        self._tn.get_socket().close()
-        self._tn = None
-
-    def _parse_state(self, output):
-        """
-        Parses the output of the state command.
-
-        :param output: list of lines that the state command printed as
-                       its output
-        :type output: list
-        """
-        for line in output:
-            status_step = ''
-            stripped = line.strip()
-            if stripped == "END":
-                continue
-            parts = stripped.split(",")
-            if len(parts) < 5:
-                continue
-            try:
-                ts, status_step, ok, ip, remote, port, _, _, _ = parts
-            except ValueError:
-                try:
-                    ts, status_step, ok, ip, remote, port, _, _ = parts
-                except ValueError:
-                    self.log.debug('Could not parse state line: %s' % line)
-
-            return status_step
-
-    def _parse_status(self, output):
-        """
-        Parses the output of the status command.
-
-        :param output: list of lines that the status command printed
-                       as its output
-        :type output: list
-        """
-        tun_tap_read = ""
-        tun_tap_write = ""
-
-        for line in output:
-            stripped = line.strip()
-            if stripped.endswith("STATISTICS") or stripped == "END":
-                continue
-            parts = stripped.split(",")
-            if len(parts) < 2:
-                continue
-
-            try:
-                text, value = parts
-            except ValueError:
-                self.log.debug('Could not parse status line %s' % line)
-                return
-            # text can be:
-            #   "TUN/TAP read bytes"
-            #   "TUN/TAP write bytes"
-            #   "TCP/UDP read bytes"
-            #   "TCP/UDP write bytes"
-            #   "Auth read bytes"
-
-            if text == "TUN/TAP read bytes":
-                tun_tap_read = value  # download
-            elif text == "TUN/TAP write bytes":
-                tun_tap_write = value  # upload
-
-        return (tun_tap_read, tun_tap_write)
-
-    def get_state(self):
-        if not self.is_connected():
-            return ""
-        state = self._parse_state(self._send_command("state"))
-        return state
-
-    def get_traffic_status(self):
-        if not self.is_connected():
-            return (None, None)
-        return self._parse_status(self._send_command("status"))
 
     def terminate(self, shutdown=False):
         """
@@ -271,125 +41,120 @@ class VPNManagement(object):
         if self.is_connected():
             self._send_command("signal SIGTERM")
         if shutdown:
-            self._cleanup_tempfiles()
+            _cleanup_tempfiles()
 
-    def _cleanup_tempfiles(self):
-        """
-        Remove all temporal files we might have left behind.
 
-        Iif self.port is 'unix', we have created a temporal socket path that,
-        under normal circumstances, we should be able to delete.
-        """
-        if self._socket_port == "unix":
-            tempfolder = _first(os.path.split(self._host))
-            if tempfolder and os.path.isdir(tempfolder):
-                try:
-                    shutil.rmtree(tempfolder)
-                except OSError:
-                    self.log.error(
-                        'Could not delete tmpfolder %s' % tempfolder)
-
-    def get_openvpn_process(self):
-        """
-        Looks for openvpn instances running.
+# TODO -- finish porting ----------------------------------------------------
 
-        :rtype: process
-        """
-        openvpn_process = None
-        for p in psutil.process_iter():
+def _cleanup_tempfiles(self):
+    """
+    Remove all temporal files we might have left behind.
+
+    Iif self.port is 'unix', we have created a temporal socket path that,
+    under normal circumstances, we should be able to delete.
+    """
+    if self._socket_port == "unix":
+        tempfolder = _first(os.path.split(self._host))
+        if tempfolder and os.path.isdir(tempfolder):
             try:
-                # XXX Not exact!
-                # Will give false positives.
-                # we should check that cmdline BEGINS
-                # with openvpn or with our wrapper
-                # (pkexec / osascript / whatever)
-
-                # This needs more work, see #3268, but for the moment
-                # we need to be able to filter out arguments in the form
-                # --openvpn-foo, since otherwise we are shooting ourselves
-                # in the feet.
-
-                if PSUTIL_2:
-                    cmdline = p.cmdline()
-                else:
-                    cmdline = p.cmdline
-                if any(map(lambda s: s.find(
-                        "LEAPOPENVPN") != -1, cmdline)):
-                    openvpn_process = p
-                    break
-            except psutil_AccessDenied:
-                pass
-        return openvpn_process
-
-    def stop_if_already_running(self):
-        """
-        Checks if VPN is already running and tries to stop it.
+                shutil.rmtree(tempfolder)
+            except OSError:
+                self.log.error(
+                    'Could not delete tmpfolder %s' % tempfolder)
 
-        Might raise OpenVPNAlreadyRunning.
+def _get_openvpn_process():
+    """
+    Looks for openvpn instances running.
 
-        :return: True if stopped, False otherwise
+    :rtype: process
+    """
+    openvpn_process = None
+    for p in psutil.process_iter():
+        try:
+            # XXX Not exact!
+            # Will give false positives.
+            # we should check that cmdline BEGINS
+            # with openvpn or with our wrapper
+            # (pkexec / osascript / whatever)
+
+            # This needs more work, see #3268, but for the moment
+            # we need to be able to filter out arguments in the form
+            # --openvpn-foo, since otherwise we are shooting ourselves
+            # in the feet.
+
+            if PSUTIL_2:
+                cmdline = p.cmdline()
+            else:
+                cmdline = p.cmdline
+            if any(map(lambda s: s.find(
+                    "LEAPOPENVPN") != -1, cmdline)):
+                openvpn_process = p
+                break
+        except psutil_AccessDenied:
+            pass
+    return openvpn_process
+
+def _stop_if_already_running():
+    """
+    Checks if VPN is already running and tries to stop it.
 
-        """
-        process = self.get_openvpn_process()
-        if not process:
-            self.log.debug('Could not find openvpn process while '
-                           'trying to stop it.')
-            return
+    Might raise OpenVPNAlreadyRunning.
 
-        self.log.debug('OpenVPN is already running, trying to stop it...')
-        cmdline = process.cmdline
+    :return: True if stopped, False otherwise
 
-        manag_flag = "--management"
+    """
+    process = _get_openvpn_process()
+    if not process:
+        self.log.debug('Could not find openvpn process while '
+                       'trying to stop it.')
+        return
 
-        if isinstance(cmdline, list) and manag_flag in cmdline:
+    log.debug('OpenVPN is already running, trying to stop it...')
+    cmdline = process.cmdline
 
-            # we know that our invocation has this distinctive fragment, so
-            # we use this fingerprint to tell other invocations apart.
-            # this might break if we change the configuration path in the
-            # launchers
+    manag_flag = "--management"
 
-            def smellslikeleap(s):
-                return "leap" in s and "providers" in s
+    if isinstance(cmdline, list) and manag_flag in cmdline:
 
-            if not any(map(smellslikeleap, cmdline)):
-                self.log.debug("We cannot stop this instance since we do not "
-                               "recognise it as a leap invocation.")
-                raise AlienOpenVPNAlreadyRunning
+        # we know that our invocation has this distinctive fragment, so
+        # we use this fingerprint to tell other invocations apart.
+        # this might break if we change the configuration path in the
+        # launchers
 
-            try:
-                index = cmdline.index(manag_flag)
-                host = cmdline[index + 1]
-                port = cmdline[index + 2]
-                self.log.debug("Trying to connect to %s:%s"
-                               % (host, port))
-                self._connect()
-
-                # XXX this has a problem with connections to different
-                # remotes. So the reconnection will only work when we are
-                # terminating instances left running for the same provider.
-                # If we are killing an openvpn instance configured for another
-                # provider, we will get:
-                # TLS Error: local/remote TLS keys are out of sync
-                # However, that should be a rare case right now.
-                self._send_command("signal SIGTERM")
-                self._close_management_socket(announce=True)
-            except (Exception, AssertionError):
-                self.log.failure('Problem trying to terminate OpenVPN')
-        else:
-            self.log.debug('Could not find the expected openvpn command line.')
-
-        process = self.get_openvpn_process()
-        if process is None:
-            self.log.debug('Successfully finished already running '
-                           'openvpn process.')
-            return True
-        else:
-            self.log.warn('Unable to terminate OpenVPN')
-            raise OpenVPNAlreadyRunning
-
-
-def _first(things):
-    try:
-        return things[0]
-    except (IndexError, TypeError):
-        return None
+        def smellslikeleap(s):
+            return "leap" in s and "providers" in s
+
+        if not any(map(smellslikeleap, cmdline)):
+            self.log.debug("We cannot stop this instance since we do not "
+                           "recognise it as a leap invocation.")
+            raise AlienOpenVPNAlreadyRunning
+
+        try:
+            index = cmdline.index(manag_flag)
+            host = cmdline[index + 1]
+            port = cmdline[index + 2]
+            self.log.debug("Trying to connect to %s:%s"
+                           % (host, port))
+            _connect()
+
+            # XXX this has a problem with connections to different
+            # remotes. So the reconnection will only work when we are
+            # terminating instances left running for the same provider.
+            # If we are killing an openvpn instance configured for another
+            # provider, we will get:
+            # TLS Error: local/remote TLS keys are out of sync
+            # However, that should be a rare case right now.
+            self._send_command("signal SIGTERM")
+        except (Exception, AssertionError):
+            log.failure('Problem trying to terminate OpenVPN')
+    else:
+        log.debug('Could not find the expected openvpn command line.')
+
+    process = _get_openvpn_process()
+    if process is None:
+        self.log.debug('Successfully finished already running '
+                       'openvpn process.')
+        return True
+    else:
+        self.log.warn('Unable to terminate OpenVPN')
+        raise OpenVPNAlreadyRunning
diff --git a/src/leap/bitmask/vpn/_telnet.py b/src/leap/bitmask/vpn/_telnet.py
deleted file mode 100644
index cfc82ef..0000000
--- a/src/leap/bitmask/vpn/_telnet.py
+++ /dev/null
@@ -1,60 +0,0 @@
-# -*- coding: utf-8 -*-
-# _telnet.py
-# Copyright (C) 2013-2017 LEAP
-#
-# This program is free software: you can redistribute it and/or modify
-# it under the terms of the GNU General Public License as published by
-# the Free Software Foundation, either version 3 of the License, or
-# (at your option) any later version.
-#
-# This program is distributed in the hope that it will be useful,
-# but WITHOUT ANY WARRANTY; without even the implied warranty of
-# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-# GNU General Public License for more details.
-#
-# You should have received a copy of the GNU General Public License
-# along with this program.  If not, see <http://www.gnu.org/licenses/>.
-
-import os
-import socket
-import telnetlib
-
-
-class ConnectionRefusedError(Exception):
-    pass
-
-
-class MissingSocketError(Exception):
-    pass
-
-
-class UDSTelnet(telnetlib.Telnet):
-    """
-    A telnet-alike class, that can listen on unix domain sockets
-    """
-
-    def open(self, host, port=23, timeout=socket._GLOBAL_DEFAULT_TIMEOUT):
-        """
-        Connect to a host. If port is 'unix', it will open a
-        connection over unix docmain sockets.
-
-        The optional second argument is the port number, which
-        defaults to the standard telnet port (23).
-        Don't try to reopen an already connected instance.
-        """
-        self.eof = 0
-        self.host = host
-        self.port = port
-        self.timeout = timeout
-
-        if self.port == "unix":
-            # unix sockets spoken
-            if not os.path.exists(self.host):
-                raise MissingSocketError()
-            self.sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
-            try:
-                self.sock.connect(self.host)
-            except socket.error:
-                raise ConnectionRefusedError()
-        else:
-            self.sock = socket.create_connection((host, port), timeout)
diff --git a/src/leap/bitmask/vpn/management.py b/src/leap/bitmask/vpn/management.py
new file mode 100644
index 0000000..b9bda6c
--- /dev/null
+++ b/src/leap/bitmask/vpn/management.py
@@ -0,0 +1,339 @@
+# -*- coding: utf-8 -*-
+# management.py
+# Copyright (c) 2012 Mike Mattice
+# Copyright (C) 2017 LEAP Encryption Access Project
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+"""
+Handles an OpenVPN process through its Management Interface.
+"""
+
+import time
+from collections import OrderedDict
+
+from twisted.internet import defer
+from twisted.protocols.basic import LineReceiver
+from twisted.internet.defer import Deferred
+from twisted.python import log
+
+from zope.interface import Interface
+
+from _human import bytes2human
+
+
+class IStateListener(Interface):
+
+    def change_state(self, state):
+        pass
+
+
+class ManagementProtocol(LineReceiver):
+
+    def __init__(self, verbose=False):
+
+        self.verbose = verbose
+        self.state = None
+        self.remote = None
+        self.rport = None
+        self.traffic = TrafficCounter()
+        self.openvpn_version = ''
+        self.pid = None
+
+        self._defs = []
+        self._statelog = OrderedDict()
+        self._linebuf = []
+        self._state_listeners = set([])
+
+    def addStateListener(self, listener):
+        """
+        A Listener must implement change_state method,
+        and it will be called with a State object.
+        """
+        self._state_listeners.add(listener)
+
+    def getStateHistory(self):
+        return self._statelog
+
+    def lineReceived(self, line):
+        if self.verbose:
+            print line
+
+        if line[0] == '>':
+            try:
+                infotype, data = line[1:].split(':', 1)
+                infotype = infotype.replace('-', '_')
+            except Exception, msg:
+                print "failed to parse '%r': %s" % (line, msg)
+                raise
+            m = getattr(self, '_handle_%s' % infotype, None)
+            if m:
+                try:
+                    m(data)
+                except Exception, msg:
+                    print "Failure in _handle_%s: %s" % (infotype, msg)
+            else:
+                self._handle_unknown(infotype, data)
+        else:
+            if line.strip() == 'END':
+                try:
+                    d = self._defs.pop(0)
+                    d.callback('\n'.join(self._linebuf))
+                except IndexError:
+                    pass
+                self._linebuf = []
+                return
+            try:
+                status, data = line.split(': ', 1)
+            except ValueError:
+                print "ERROR PARSING:", line
+                return
+            if status in ('ERROR', 'SUCCESS'):
+                try:
+                    d = self._defs.pop(0)
+                    if status == 'SUCCESS':
+                        d.callback(line)
+                    else:
+                        d.errback(line)
+                except:
+                    pass
+            else:
+                self._linebuf.append(line)
+
+    def _handle_unknown(self, infotype, data):
+        log.msg('Received unhandled infotype %s with data %s' %
+                (infotype, data))
+
+    def _handle_BYTECOUNT(self, data):
+        down, up = data.split(',')
+        self.traffic.update(down, up, time.time())
+
+    def _handle_ECHO(self, data):
+        pass
+
+    def _handle_FATAL(self, data):
+        pass
+
+    def _handle_HOLD(self, data):
+        pass
+
+    def _handle_INFO(self, data):
+        pass
+
+    def _handle_LOG(self, data):
+        pass
+
+    def _handle_NEED_OK(self, data):
+        pass
+
+    def _handle_NEED_STR(self, data):
+        pass
+
+    def _handle_STATE(self, data):
+        data = data.strip().split(',')
+        remote = rport = None
+        try:
+            if len(data) == 9:
+                (ts, state, verbose, localtun,
+                 remote, rport, laddr, lport, ip6) = data
+            elif len(data) == 8:
+                ts, state = data[:2]
+        except Exception as exc:
+            print "ERROR", exc
+            log.error('Failure parsing data: %s' % exc)
+
+        if state != self.state:
+            now = time.time()
+            stateobj = State(state, ts)
+            self._statelog[now] = stateobj
+            for listener in self._state_listeners:
+                listener.change_state(stateobj)
+        self.state = stateobj
+        self.remote = remote
+        self.rport = rport
+
+    def _pushdef(self):
+        d = Deferred()
+        self._defs.append(d)
+        return d
+
+    def byteCount(self, interval=0):
+        d = self._pushdef()
+        self.sendLine('bytecount %d' % (interval,))
+        return d
+
+    def signal(self, signal='SIGTERM'):
+        d = self._pushdef()
+        self.sendLine('signal %s' % (signal,))
+        return d
+
+    def _parseHoldstatus(self, result):
+        return result.split('=')[0] == '1'
+
+    def hold(self, p=''):
+        d = self._pushdef()
+        self.sendLine('hold %s' % (p,))
+        if p == '':
+            d.addCallback(self._parseHoldstatus)
+        return d
+
+    def _parsePid(self, result):
+        self.pid = int(result.split('=')[1])
+
+    def get_pid(self):
+        d = self._pushdef()
+        self.sendLine('pid')
+        d.addCallback(self._parsePid)
+        return d
+
+    def logOn(self):
+        d = self._pushdef()
+        self.sendLine('log on')
+        return d
+
+    def stateOn(self):
+        d = self._pushdef()
+        self.sendLine('state on')
+        return d
+
+    def _parseVersion(self, data):
+        version = data.split('\n')[0].split(':')[1]
+        self.openvpn_version = version.strip()
+
+    def getVersion(self):
+        d = self._pushdef()
+        self.sendLine('version')
+        d.addCallback(self._parseVersion)
+        return d
+
+    def getInfo(self):
+        state = self._statelog.values()[-1]
+        return {
+            'remote': self.remote,
+            'rport': self.rport,
+            'state': state.state,
+            'state_simple': state.simple,
+            'state_legend': state.legend,
+            'openvpn_version': self.openvpn_version,
+            'pid': self.pid,
+            'traffic_down_total': self.traffic.down,
+            'traffic_up_total': self.traffic.up}
+
+
+class State(object):
+
+    """
+    Possible States in an OpenVPN connection, according to the
+    OpenVPN Management documentation.
+    """
+
+    CONNECTING = 'CONNECTING'
+    WAIT = 'WAIT'
+    AUTH = 'AUTH'
+    GET_CONFIG = 'GET_CONFIG'
+    ASSIGN_IP = 'ASSIGN_IP'
+    ADD_ROUTES = 'ADD_ROUTES'
+    CONNECTED = 'CONNECTED'
+    RECONNECTING = 'RECONNECTING'
+    EXITING = 'EXITING'
+
+    OFF = 'OFF'
+    ON = 'ON'
+    STARTING = 'STARTING'
+    STOPPING = 'STOPPING'
+    FAILED = 'FAILED'
+
+    _legend = {
+        'CONNECTING': 'Connecting to remote server',
+        'WAIT': 'Waiting from initial response from server',
+        'AUTH': 'Authenticating with server',
+        'GET_CONFIG': 'Downloading configuration options from server',
+        'ASSIGN_IP': 'Assigning IP address to virtual network interface',
+        'ADD_ROUTES': 'Adding routes to system',
+        'CONNECTED': 'Initialization Sequence Completed',
+        'RECONNECTING': 'A restart has occurred',
+        'EXITING': 'A graceful exit is in progress'
+    }
+
+    _simple = {
+        'CONNECTING': STARTING,
+        'WAIT': STARTING,
+        'AUTH': STARTING,
+        'GET_CONFIG': STARTING,
+        'ASSIGN_IP': STARTING,
+        'ADD_ROUTES': STARTING,
+        'CONNECTED': ON,
+        'RECONNECTING': STARTING,
+        'EXITING': STOPPING
+    }
+
+    def __init__(self, state, timestamp):
+        self.state = state
+        self.timestamp = timestamp
+
+    @classmethod
+    def get_legend(cls, state):
+        return cls._legend.get(state)
+
+    @classmethod
+    def get_simple(cls, state):
+        return cls._simple.get(state)
+
+    @property
+    def simple(self):
+        return self.get_simple(self.state)
+
+    @property
+    def legend(self):
+        return self.get_legend(self.state)
+
+    def __repr__(self):
+        return '<State: %s [%s]>' % (
+            self.state, time.ctime(int(self.timestamp)))
+
+
+class TrafficCounter(object):
+
+    CAPACITY = 60
+
+    def __init__(self):
+        self.down = None
+        self.up = None
+        self._buf = OrderedDict()
+
+    def update(self, down, up, ts):
+        i_down = int(down)
+        i_up = int(up)
+        self.down = i_down
+        self.up = i_up
+        if len(self._buf) > self.CAPACITY:
+            self._buf.pop(self._buf.keys()[0])
+        self._buf[ts] = i_down, i_up
+
+    def get_rate(self, human=True):
+        points = self._buf.items()
+        if len(points) < 2:
+            return ['NA', 'NA']
+        ts1, prev = points[-2]
+        ts2, last = points[-1]
+        rate_down = _get_rate(last[0], prev[0], ts2, ts1)
+        rate_up = _get_rate(last[1], prev[1], ts2, ts1)
+        rates = rate_down, rate_up
+        if human:
+            rates = map(bytes2human, rates)
+        return rates
+
+
+def _get_rate(p2, p1, ts2, ts1):
+    return ((1.0 * (p2 - p1)) / (ts2 - ts1))
diff --git a/tests/unit/vpn/session1.data b/tests/unit/vpn/session1.data
new file mode 100644
index 0000000..5381bad
--- /dev/null
+++ b/tests/unit/vpn/session1.data
@@ -0,0 +1,34 @@
+>INFO:OpenVPN Management Interface Version 1 -- type 'help' for more info
+SUCCESS: real-time log notification set to ON
+OpenVPN Version: OpenVPN 2.4.0 x86_64-pc-linux-gnu [SSL (OpenSSL)] [LZO] [LZ4] [EPOLL] [PKCS11] [MH/PKTINFO] [AEAD] built on Jun 22 2017
+Management Version: 1
+END
+SUCCESS: pid=23783
+SUCCESS: real-time state notification set to ON
+SUCCESS: bytecount interval changed
+>BYTECOUNT:26,14
+>STATE:1503010298,AUTH,,,,,,
+>LOG:1503010302,I,[otter.demo.bitmask.net] Peer Connection Initiated with [AF_INET]46.165.242.169:443
+>STATE:1503010303,GET_CONFIG,,,,,,
+>BYTECOUNT:5573,3716
+>LOG:1503010303,W,Note: option tun-ipv6 is ignored because modern operating systems do not need special IPv6 tun handling anymore.
+>LOG:1503010303,I,TUN/TAP device tun0 opened
+>LOG:1503010303,D,do_ifconfig, tt->did_ifconfig_ipv6_setup=1
+>STATE:1503010303,ASSIGN_IP,,10.42.0.18,,,,,2001:db8:123::1010
+>LOG:1503010303,I,/sbin/ip link set dev tun0 up mtu 1500
+>LOG:1503010303,I,/sbin/ip addr add dev tun0 10.42.0.18/21 broadcast 10.42.7.255
+>LOG:1503010303,I,/sbin/ip -6 addr add 2001:db8:123::1010/64 dev tun0
+>LOG:1503010303,W,ERROR: Linux route add command failed: external program exited with error status: 2
+>LOG:1503010304,I,add_route_ipv6(2000::/3 -> 2001:db8:123::1 metric -1) dev tun0
+>LOG:1503010304,I,GID set to nogroup
+>LOG:1503010304,I,UID set to nobody
+>LOG:1503010304,I,Initialization Sequence Completed
+>STATE:1503010304,CONNECTED,SUCCESS,10.42.0.18,46.165.242.169,443,,,2001:db8:123::1010
+>BYTECOUNT:6279,4102
+>BYTECOUNT:6577,4400
+>BYTECOUNT:6923,4847
+>BYTECOUNT:7237,4996
+>BYTECOUNT:7732,5443
+>BYTECOUNT:10355,6055
+>BYTECOUNT:17168,7657
+
diff --git a/tests/unit/vpn/session2.data b/tests/unit/vpn/session2.data
new file mode 100644
index 0000000..cd2ece5
--- /dev/null
+++ b/tests/unit/vpn/session2.data
@@ -0,0 +1,38 @@
+>INFO:OpenVPN Management Interface Version 1 -- type 'help' for more info
+SUCCESS: real-time log notification set to ON
+OpenVPN Version: OpenVPN 2.4.0 x86_64-pc-linux-gnu [SSL (OpenSSL)] [LZO] [LZ4] [EPOLL] [PKCS11] [MH/PKTINFO] [AEAD] built on Jun 22 2017
+Management Version: 1
+END
+SUCCESS: pid=30641
+SUCCESS: real-time state notification set to ON
+SUCCESS: bytecount interval changed
+>BYTECOUNT:26,14
+>STATE:1503079741,AUTH,,,,,,
+>BYTECOUNT:5168,3078
+>LOG:1503079746,I,[otter.demo.bitmask.net] Peer Connection Initiated with [AF_INET]46.165.242.169:443
+>STATE:1503079747,GET_CONFIG,,,,,,
+>LOG:1503079748,W,Note: option tun-ipv6 is ignored because modern operating systems do not need special IPv6 tun handling anymore.
+>LOG:1503079748,I,TUN/TAP device tun0 opened
+>LOG:1503079748,D,do_ifconfig, tt->did_ifconfig_ipv6_setup=1
+>STATE:1503079748,ASSIGN_IP,,10.42.0.12,,,,,2001:db8:123::100a
+>LOG:1503079748,I,/sbin/ip link set dev tun0 up mtu 1500
+>LOG:1503079748,I,/sbin/ip addr add dev tun0 10.42.0.12/21 broadcast 10.42.7.255
+>LOG:1503079748,I,/sbin/ip -6 addr add 2001:db8:123::100a/64 dev tun0
+>LOG:1503079748,W,ERROR: Linux route add command failed: external program exited with error status: 2
+>LOG:1503079748,I,add_route_ipv6(2000::/3 -> 2001:db8:123::1 metric -1) dev tun0
+>LOG:1503079748,I,GID set to nogroup
+>LOG:1503079748,I,UID set to nobody
+>LOG:1503079748,I,Initialization Sequence Completed
+>STATE:1503079748,CONNECTED,SUCCESS,10.42.0.12,46.165.242.169,443,,,2001:db8:123::100a
+SUCCESS: signal SIGTERM thrown
+>LOG:1503079751,W,ERROR: Linux route delete command failed: external program exited with error status: 2
+>LOG:1503079751,W,ERROR: Linux route delete command failed: external program exited with error status: 2
+>LOG:1503079751,W,ERROR: Linux route delete command failed: external program exited with error status: 2
+>LOG:1503079751,I,delete_route_ipv6(2000::/3)
+>LOG:1503079751,W,ERROR: Linux route -6/-A inet6 del command failed: external program exited with error status: 2
+>LOG:1503079751,I,/sbin/ip addr del dev tun0 10.42.0.12/21
+>LOG:1503079751,W,Linux ip addr del failed: external program exited with error status: 2
+>LOG:1503079751,I,/sbin/ip -6 addr del 2001:db8:123::100a/64 dev tun0
+>LOG:1503079751,W,Linux ip -6 addr del failed: external program exited with error status: 2
+>LOG:1503079751,I,SIGTERM[hard,] received, process exiting
+>STATE:1503079751,EXITING,SIGTERM,,,,,
diff --git a/tests/unit/vpn/test_management.py b/tests/unit/vpn/test_management.py
new file mode 100644
index 0000000..c70d768
--- /dev/null
+++ b/tests/unit/vpn/test_management.py
@@ -0,0 +1,129 @@
+# -*- coding: utf-8 -*-
+# test_management.py
+# Copyright (C) 2017 LEAP Encryption Access Project
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+"""
+Tests for the VPN Management Interface
+"""
+
+import StringIO
+
+from twisted.trial import unittest
+from leap.bitmask.vpn.management import ManagementProtocol
+
+
+session1 = open('session1.data').readlines()
+session2 = open('session2.data').readlines()
+
+
+def feed_the_protocol(protocol, data):
+    for line in data:
+        protocol.lineReceived(line)
+
+
+class StateListener(object):
+
+    def __init__(self):
+        self.states = []
+
+    def change_state(self, state):
+        self.states.append(state)
+
+
+class ManagementTestCase(unittest.TestCase):
+
+    def test_final_state_is_connected(self):
+        proto = ManagementProtocol()
+        feed_the_protocol(proto, session1)
+        assert proto.state.state == 'CONNECTED'
+        assert proto.state.simple == 'ON'
+        assert proto.remote == '46.165.242.169'
+
+    def test_final_state_stopping(self):
+        proto = ManagementProtocol()
+        feed_the_protocol(proto, session2)
+        assert proto.state.state == 'EXITING'
+        assert proto.state.simple == 'STOPPING'
+
+    def test_get_state_history(self):
+        proto = ManagementProtocol()
+        feed_the_protocol(proto, session1)
+        log = proto.getStateHistory()
+        states = [st.state for st in log.values()]
+        assert len(log) == 4
+        assert states == ['AUTH', 'GET_CONFIG', 'ASSIGN_IP', 'CONNECTED']
+
+    def test_state_listener(self):
+        proto = ManagementProtocol()
+        listener = StateListener()
+        proto.addStateListener(listener)
+        feed_the_protocol(proto, session1)
+        states = [st.state for st in listener.states]
+        assert states == ['AUTH', 'GET_CONFIG', 'ASSIGN_IP', 'CONNECTED']
+
+    def test_bytecount(self):
+        proto = ManagementProtocol()
+        feed_the_protocol(proto, session1)
+        assert proto.traffic.down == 17168
+        assert proto.traffic.up == 7657
+
+    def test_bytecount_rate(self):
+        proto = ManagementProtocol()
+        proto.traffic.update(1024, 512, 1)
+        proto.traffic.update(2048, 1024, 2)
+        print proto.traffic._buf
+        assert proto.traffic.down == 2048
+        assert proto.traffic.up == 1024
+        assert proto.traffic.get_rate() == ['1.0 K', '512.0 B']
+
+    def test_get_pid(self):
+        proto = ManagementProtocol()
+        proto.transport = StringIO.StringIO()
+        assert proto.pid == None
+        proto.get_pid()
+        pid_lines = ['SUCCESS: pid=99999']
+        feed_the_protocol(proto, pid_lines)
+        assert proto.pid == 99999
+
+    def test_parse_version_string(self):
+        proto = ManagementProtocol()
+        proto.transport = StringIO.StringIO()
+        assert proto.openvpn_version == ''
+        feed_the_protocol(proto, session1[2:4])
+        proto.getVersion()
+        feed_the_protocol(proto, ['END'])
+        assert proto.openvpn_version.startswith('OpenVPN 2.4.0')
+
+    def test_get_info(self):
+        proto = ManagementProtocol()
+        proto.transport = StringIO.StringIO()
+        feed_the_protocol(proto, session1)
+
+        feed_the_protocol(proto, session1[2:4])
+        proto.getVersion()
+        feed_the_protocol(proto, ['END'])
+        proto.get_pid()
+        pid_lines = ['SUCCESS: pid=23783']
+        feed_the_protocol(proto, pid_lines)
+
+        info = proto.getInfo()
+        assert info['remote'] == '46.165.242.169'
+        assert info['rport'] == '443'
+        assert info['state'] == 'CONNECTED'
+        assert info['state_simple'] == 'ON'
+        assert info['state_legend'] == 'Initialization Sequence Completed'
+        assert info['openvpn_version'].startswith('OpenVPN 2.4.0')
+        assert info['pid'] == 23783