add cleanup code for vnc connections

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

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):

requirements.txt

@@ -12,3 +12,4 @@ raven
 paramiko
 websockify
 sshtunnel
+psutil