diff --git a/HACKING/build_env.sh b/HACKING/build_env.sh new file mode 100755 index 0000000..0590705 --- /dev/null +++ b/HACKING/build_env.sh @@ -0,0 +1,2 @@ +#!/bin/bash +podman build . --tag=proxstar diff --git a/proxstar/__init__.py b/proxstar/__init__.py index e1b889b..33e43fe 100644 --- a/proxstar/__init__.py +++ b/proxstar/__init__.py @@ -252,6 +252,9 @@ def vm_power(vmid, action): connect_proxmox() if user.rtp or int(vmid) in user.allowed_vms: vm = VM(vmid) + vnc_token_key = f'vnc_token|{vmid}' + vnc_token = redis_conn.get(vnc_token_key).decode('utf-8') # For deleting the token from redis later + print(f'vnc_token = {vnc_token}') if action == 'start': vmconfig = vm.config usage_check = user.check_usage(vmconfig['cores'], vmconfig['memory'], 0) @@ -260,16 +263,18 @@ def vm_power(vmid, action): vm.start() elif action == 'stop': vm.stop() - # TODO (willnilges): Replace with remove target function or something - # send_stop_ssh_tunnel(vmid) + delete_vnc_target(token=vnc_token) + redis_conn.delete(vnc_token_key) elif action == 'shutdown': vm.shutdown() - # send_stop_ssh_tunnel(vmid) + delete_vnc_target(token=vnc_token) + redis_conn.delete(vnc_token_key) elif action == 'reset': vm.reset() elif action == 'suspend': vm.suspend() - # send_stop_ssh_tunnel(vmid) + delete_vnc_target(token=vnc_token) + redis_conn.delete(vnc_token_key) elif action == 'resume': vm.resume() return '', 200 @@ -290,6 +295,7 @@ def vm_console(vmid): ) node = f'{vm.node}.csh.rit.edu' token = add_vnc_target(node, vnc_port) + redis_conn.set(f'vnc_token|{vmid}', str(token)) # Store the VNC token in Redis. return { 'host': app.config['VNC_HOST'], 'port': app.config['VNC_PORT'], diff --git a/proxstar/tasks.py b/proxstar/tasks.py index 0e928a8..0063e6e 100644 --- a/proxstar/tasks.py +++ b/proxstar/tasks.py @@ -6,6 +6,7 @@ import psycopg2 import requests from flask import Flask from rq import get_current_job +from redis import Redis from sqlalchemy import create_engine from sqlalchemy.orm import sessionmaker @@ -32,7 +33,6 @@ else: config = os.path.join(app.config.get('ROOT_DIR', os.getcwd()), 'config.py') app.config.from_pyfile(config) - def connect_db(): engine = create_engine(app.config['SQLALCHEMY_DATABASE_URI']) Base.metadata.bind = engine diff --git a/proxstar/vnc.py b/proxstar/vnc.py index 536f39d..1e92856 100644 --- a/proxstar/vnc.py +++ b/proxstar/vnc.py @@ -56,15 +56,22 @@ def add_vnc_target(node, port): return token -def delete_vnc_target(node, port): +def delete_vnc_target(node=None, port=None, token=None): targets = get_vnc_targets() - target = next((target for target in targets if target['host'] == f'{node}:{port}'), None) + if node is not None and port is not None: + target = next((target for target in targets if target['host'] == f'{node}:{port}'), None) + elif token is not None: + target = next((target for target in targets if target['token'] == f'{token}'), None) + else: + raise ValueError("Need either a node and port, or a token.") if target: targets.remove(target) target_file = open(app.config['WEBSOCKIFY_TARGET_FILE'], 'w') for target in targets: target_file.write(f"{target['token']}: {target['host']}\n") target_file.close() + else: + raise LookupError("Target does not exist") def open_vnc_session(vmid, node, proxmox_user, proxmox_pass):