From 115909b5cde429c4e8872aad9f31b9af0ff1474a Mon Sep 17 00:00:00 2001
From: Jordan Rodgers <com6056@gmail.com>
Date: Tue, 20 Feb 2018 01:10:42 -0500
Subject: [PATCH] add cleanup code for vnc connections

---
 proxstar/__init__.py | 20 +++++++++++++++++---
 proxstar/vnc.py      | 12 ++++++++++--
 requirements.txt     |  1 +
 3 files changed, 28 insertions(+), 5 deletions(-)

diff --git a/proxstar/__init__.py b/proxstar/__init__.py
index 588c2d4..b8b1ed7 100644
--- a/proxstar/__init__.py
+++ b/proxstar/__init__.py
@@ -1,5 +1,6 @@
 import os
 import time
+import psutil
 import atexit
 import subprocess
 from rq import Queue
@@ -456,11 +457,24 @@ def allowed_users(user):
         return '', 403
 
 
-@app.route("/targets/clear")
+@app.route("/vnc/cleanup")
 @auth.oidc_auth
-def clear_targets():
+def cleanup_vnc():
     if 'rtp' in session['userinfo']['groups']:
-        clear_vnc_targets()
+        for target in get_vnc_targets():
+            tunnel = next((tunnel for tunnel in ssh_tunnels
+                           if tunnel.local_bind_port == int(target['port'])),
+                          None)
+            if tunnel:
+                if not next((conn for conn in psutil.net_connections()
+                             if conn.laddr[1] == int(target['port'])
+                             and conn.status == 'ESTABLISHED'), None):
+                    try:
+                        tunnel.stop()
+                    except:
+                        pass
+                    ssh_tunnels.remove(tunnel)
+                    delete_vnc_target(target['port'])
         return '', 200
     else:
         return '', 403
diff --git a/proxstar/vnc.py b/proxstar/vnc.py
index aacf6e9..ad2dd55 100644
--- a/proxstar/vnc.py
+++ b/proxstar/vnc.py
@@ -60,8 +60,16 @@ def add_vnc_target(port):
         return token
 
 
-def clear_vnc_targets():
-    open(app.config['WEBSOCKIFY_TARGET_FILE'], 'w').close()
+def delete_vnc_target(port):
+    targets = get_vnc_targets()
+    target = next((target for target in targets if target['port'] == port),
+                  None)
+    if target:
+        targets.remove(target)
+        target_file = open(app.config['WEBSOCKIFY_TARGET_FILE'], 'w')
+        for target in targets:
+            target_file.write("{}\n".format(target))
+        target_file.close()
 
 
 def start_ssh_tunnel(node, port):
diff --git a/requirements.txt b/requirements.txt
index 92c6c5d..2346181 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -12,3 +12,4 @@ raven
 paramiko
 websockify
 sshtunnel
+psutil