From e8902e50e9eb9f468346f36436811f13aee5de9b Mon Sep 17 00:00:00 2001
From: mantaohuang <mantao@mit.edu>
Date: Sun, 5 Apr 2020 18:27:38 -0400
Subject: [PATCH] add monitoring and speed display

---
 openvpn.py | 81 ++++++++++++++++++++++++++++++++++++++++++++++++++----
 test.py    |  5 ++++
 2 files changed, 81 insertions(+), 5 deletions(-)

diff --git a/openvpn.py b/openvpn.py
index 493ec8a..5ac06e4 100644
--- a/openvpn.py
+++ b/openvpn.py
@@ -7,6 +7,10 @@ import os
 import time
 import openvpn_api
 import signal
+import psutil
+import subprocess
+from datetime import datetime
+import humanize
 
 
 def generate_config(in_fp, cfg):
@@ -45,10 +49,14 @@ class Openvpn:
         self.proc = None
         self.exit_future = None
         self.additional_cfg = additional_cfg
-        self.run_task = None
+        self.run_task = []
         self.openvpn_api = None
         self.name = name
         self.pid_fp = None
+        # TODO: update paths function
+        self.io_stat_fp = os.path.join(self.folder_path, "io_stat.txt")
+        self.ping_stat_fp = os.path.join(self.folder_path, "ping_stat.txt")
+
         if loop:
             self.loop = loop
         else:
@@ -78,8 +86,10 @@ class Openvpn:
             self.status = RUNNING
             config_fp = self.generate_config_file()
             route_up_script = os.path.join(self.folder_path, "test.sh")
-            cmd = " ".join(["openvpn", "--config", config_fp, "--route-noexec", "--route-up", route_up_script ,"--script-security","2"])
-            self.run_task = self.loop.create_task(self.run(cmd))
+            cmd = " ".join(["openvpn", "--config", config_fp, "--route-noexec",
+                            "--route-up", route_up_script, "--script-security", "2"])
+            self.run_task.append(self.loop.create_task(self.run(cmd)))
+            self.run_task.append(self.loop.create_task(self.monitor_task()))
 
     def get_log(self):
         # regenerate log_fp
@@ -101,7 +111,9 @@ class Openvpn:
                 print("kill failed:", err)
 
             self.status = IDLE
-            self.run_task.cancel()
+            for task in self.run_task:
+                task.cancel()
+            self.run_task = []
             try:
                 os.remove(self.log_fp)
             except Exception as err:
@@ -111,6 +123,65 @@ class Openvpn:
         await self.stop()
         self.start()
 
+    def get_io_stat(self):
+        try:
+            return open(self.io_stat_fp, "r").read()
+        except:
+            return ""
+
+    def get_ping_stat(self):
+        try:
+            return open(self.ping_stat_fp, "r").read()
+        except:
+            return ""
+
+    def get_latest_speed(self):
+        try:
+            buf = open(self.io_stat_fp, "r").readlines()
+            buf = [[float(i) for i in l.split(",")] for l in buf[-2:]]
+            assert len(buf) >= 2, "not enough points, need >=2"
+            delta_t = buf[1][0] - buf[0][0]
+            down_speed = (buf[1][1] - buf[0][1])/delta_t
+            up_speed = (buf[1][2] - buf[0][2])/delta_t
+
+            down_speed = humanize.naturalsize(down_speed)
+            up_speed = humanize.naturalsize(up_speed)
+
+            return f"Down: {down_speed}/s, Up: {up_speed}/s"
+        except:
+            return "no data"
+
+    def monitor_action(self):
+        """
+        monitor the I/O statistics and connectivity
+        TODO: seperate module for monitoring
+        """
+        io_count = psutil.net_io_counters(pernic=True)
+        assert self.interface in io_count, "cannot find interface, nic is not there"
+        io_count = io_count[self.interface]
+        open(self.io_stat_fp, "a").write(
+            f"{datetime.utcnow().timestamp()},{io_count.bytes_recv},{io_count.bytes_recv}\n")
+        #command = ["ping", "-w", "1", "-c", "1", "8.8.8.8"]
+        command = ["fping", "-C", "1", "-t200", "-q", "8.8.8.8"]
+        p = subprocess.Popen(
+            command, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+        result = p.communicate()
+        if p.returncode == 0:
+            result = result[1].decode().strip().split(':')[-1]
+        else:
+            result = "-1"
+        open(self.ping_stat_fp, "a").write(
+            f"{datetime.utcnow().timestamp()},{result}\n")
+
+    async def monitor_task(self):
+        print("running monitor task")
+        while self.status == RUNNING:
+            try:
+                self.monitor_action()
+            except Exception as e:
+                print("monitoring error:", e)
+            await asyncio.sleep(5)
+
     async def run(self, cmd, group="openvpn"):
         print(f"run: {cmd}")
         self.exit_future = asyncio.Future()
@@ -162,4 +233,4 @@ if __name__ == "__main__":
     cfg_fp = "/home/mantao/Desktop/t/TCP_Files/UK2-TCP.ovpn"
     o1 = Openvpn(cfg_fp, "tun0", folder_fp, 8001, additional_cfg={
         "auth-user-pass": "/home/mantao/Desktop/t/fast.txt"})
-    o1.start()
+    # o1.start()
diff --git a/test.py b/test.py
index c3cefbc..fb509d2 100644
--- a/test.py
+++ b/test.py
@@ -24,6 +24,11 @@ class MainHandler(tornado.web.RequestHandler):
                 buf += "log: <br/>\n"
                 buf += "<pre>\n"
                 buf += op.get_log()
+                buf += f"\nSpeed: {op.get_latest_speed()}\n"
+                buf += "\nio_stat\n"
+                buf += "\n".join(op.get_io_stat().splitlines()[-5])
+                buf += "\nping_stat\n"
+                buf += "\n".join(op.get_ping_stat().splitlines()[-5])
                 buf += "\n</pre>\n"
         self.write(buf)