From 7ab74fb2e47ce752128ee6e5c71ea46c6a3f1a89 Mon Sep 17 00:00:00 2001 From: Will Nilges Date: Sun, 10 Jul 2022 17:43:14 -0400 Subject: [PATCH] Update websockify-related tasks --- HACKING/launch_env.sh | 2 +- HACKING/stop_env.sh | 6 ++--- proxstar/__init__.py | 53 +++++++++++++++++++++++-------------------- proxstar/tasks.py | 8 ++++++- proxstar/vnc.py | 4 +++- 5 files changed, 43 insertions(+), 30 deletions(-) diff --git a/HACKING/launch_env.sh b/HACKING/launch_env.sh index a22a1b6..d17f027 100755 --- a/HACKING/launch_env.sh +++ b/HACKING/launch_env.sh @@ -3,4 +3,4 @@ podman run --rm -d --network=proxstar --name=proxstar-redis redis:alpine podman run --rm -d --network=proxstar --name=proxstar-postgres -e POSTGRES_PASSWORD=changeme -v ./HACKING/proxstar-postgres/volume:/var/lib/postgresql/data:Z proxstar-postgres podman run --rm -d --network=proxstar --name=proxstar-rq-scheduler --env-file=HACKING/.env --entrypoint ./start_scheduler.sh proxstar podman run --rm -d --network=proxstar --name=proxstar-rq --env-file=HACKING/.env --entrypoint ./start_worker.sh proxstar -podman run --rm -d --network=proxstar --name=proxstar -p 8000:8000 --env-file=HACKING/.env --entrypoint='["python3", "wsgi.py"]' proxstar +podman run --rm -it --network=proxstar --name=proxstar -p 8000:8000 -p 8081:8081 --env-file=HACKING/.env --entrypoint='["python3", "wsgi.py"]' proxstar \ No newline at end of file diff --git a/HACKING/stop_env.sh b/HACKING/stop_env.sh index 4e29fa2..304e996 100755 --- a/HACKING/stop_env.sh +++ b/HACKING/stop_env.sh @@ -1,6 +1,6 @@ #!/bin/bash -podman stop proxstar -podman stop proxstar-rq -podman stop proxstar-rq-scheduler +podman kill proxstar +podman kill proxstar-rq +podman kill proxstar-rq-scheduler podman stop proxstar-redis podman stop proxstar-postgres diff --git a/proxstar/__init__.py b/proxstar/__init__.py index 83d2f6a..0843777 100644 --- a/proxstar/__init__.py +++ b/proxstar/__init__.py @@ -125,6 +125,7 @@ if 'process_expiring_vms' not in scheduler: logging.info('adding process expiring VMs task to scheduler') scheduler.cron('0 5 * * *', id='process_expiring_vms', func=process_expiring_vms_task) +# FIXME (willnilges): is this operating in the right container? if 'cleanup_vnc' not in scheduler: logging.info('adding cleanup VNC task to scheduler') scheduler.schedule( @@ -134,7 +135,6 @@ if 'cleanup_vnc' not in scheduler: interval=3600, ) - def add_rq_dashboard_auth(blueprint): @blueprint.before_request @auth.oidc_auth @@ -575,29 +575,34 @@ def allowed_users(user): @app.route('/console/cleanup', methods=['POST']) def cleanup_vnc(): if request.form['token'] == app.config['VNC_CLEANUP_TOKEN']: - 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 + with open(app.config['WEBSOCKIFY_TARGET_FILE'], 'w') as targets: + targets.truncate() + return '', 200 + return '', 403 + # if request.form['token'] == app.config['VNC_CLEANUP_TOKEN']: + # 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 @app.route('/template//disk') diff --git a/proxstar/tasks.py b/proxstar/tasks.py index bf16788..311d59f 100644 --- a/proxstar/tasks.py +++ b/proxstar/tasks.py @@ -227,8 +227,14 @@ def setup_template_task(template_id, name, user, ssh_key, cores, memory): def cleanup_vnc_task(): + """Removes all open VNC sessions. This runs in the RQ worker, and so + needs to be routed properly via the Proxstar API + TODO (willnilges): Use API, track the task IDs, and kill only the finished + ones every couple of minutes + https://github.com/ComputerScienceHouse/proxstar/issues/153 + """ requests.post( - 'https://{}/console/cleanup'.format(app.config['SERVER_NAME']), + f'https://{app.config["VNC_HOST"]}/console/cleanup', data={'token': app.config['VNC_CLEANUP_TOKEN']}, verify=False, ) diff --git a/proxstar/vnc.py b/proxstar/vnc.py index a1dc793..57e5898 100644 --- a/proxstar/vnc.py +++ b/proxstar/vnc.py @@ -32,16 +32,18 @@ def get_vnc_targets(): target_dict = {} values = line.strip().split(':') target_dict['token'] = values[0] - target_dict['host'] = values[1] + values[2] + target_dict['host'] = f'{values[1].strip()}:{values[2]}' targets.append(target_dict) target_file.close() return targets def add_vnc_target(node, port): # TODO (willnilges): This doesn't throw an error if the target file is wrong. + # TODO (willnilges): This will duplicate targets targets = get_vnc_targets() target = next((target for target in targets if target['host'] == f'{node}:{port}'), None) if target: + print('Host is already in the targets file') return target['token'] else: target_file = open(app.config['WEBSOCKIFY_TARGET_FILE'], 'a')