Use Redis to link VNC tokens to VMID

Author: WillNilges

HACKING/build_env.sh

@@ -0,0 +1,2 @@
+#!/bin/bash
+podman build . --tag=proxstar 

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'],

proxstar/tasks.py

@@ -6,6 +6,7 @@
 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 @@
     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

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