From 115909b5cde429c4e8872aad9f31b9af0ff1474a Mon Sep 17 00:00:00 2001 From: Jordan Rodgers 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