[feature] add twisted protocol for handling openvpn management

1 files changed, 129 insertions, 0 deletions

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