From 576c6c13385986857ce8a013e0b6db1b44f7b8f0 Mon Sep 17 00:00:00 2001
From: Jordan Rodgers <com6056@gmail.com>
Date: Mon, 12 Feb 2018 20:55:09 -0500
Subject: [PATCH] beginning of vnc work, hopefully some of this works

---
 .gitmodules                        |  3 ++
 config.py                          |  5 +++
 proxstar/__init__.py               | 51 ++++++++++++++++-----
 proxstar/proxmox.py                |  6 ---
 proxstar/static/css/styles.css     |  7 ---
 proxstar/static/js/script.js       | 20 +++++++++
 proxstar/static/noVNC              |  1 +
 proxstar/templates/vm_details.html |  4 +-
 proxstar/vnc.py                    | 71 ++++++++++++++++++++++++++++++
 requirements.txt                   |  1 +
 10 files changed, 143 insertions(+), 26 deletions(-)
 create mode 160000 proxstar/static/noVNC
 create mode 100644 proxstar/vnc.py

diff --git a/.gitmodules b/.gitmodules
index a09ca09..dba2cdd 100644
--- a/.gitmodules
+++ b/.gitmodules
@@ -1,3 +1,6 @@
 [submodule "proxstar/static/css/csh-material-bootstrap"]
 	path = proxstar/static/css/csh-material-bootstrap
 	url = https://github.com/ComputerScienceHouse/csh-material-bootstrap.git
+[submodule "proxstar/static/noVNC"]
+	path = proxstar/static/noVNC
+	url = https://github.com/novnc/noVNC.git
diff --git a/config.py b/config.py
index 367d036..b1f24c9 100644
--- a/config.py
+++ b/config.py
@@ -31,6 +31,7 @@ PROXMOX_HOSTS = [
 PROXMOX_USER = environ.get('PROXSTAR_PROXMOX_USER', '')
 PROXMOX_PASS = environ.get('PROXSTAR_PROXMOX_PASS', '')
 PROXMOX_ISO_STORAGE = environ.get('PROXSTAR_PROXMOX_ISO_STORAGE', 'nfs-iso')
+PROXMOX_SSH_KEY = environ.get('PROXSTAR_PROXMOX_SSH_KEY', '')
 
 # STARRS
 STARRS_DB_HOST = environ.get('PROXSTAR_STARRS_DB_HOST', '')
@@ -50,3 +51,7 @@ SQLALCHEMY_DATABASE_URI = environ.get('PROXSTAR_SQLALCHEMY_DATABASE_URI', '')
 # REDIS
 REDIS_HOST = environ.get('PROXSTAR_REDIS_HOST', 'localhost')
 REDIS_PORT = int(environ.get('PROXSTAR_REDIS_PORT', '6379'))
+
+# VNC
+WEBSOCKIFY_PATH = environ.get('PROXSTAR_WEBSOCKIFY_PATH', '/opt/app-root/bin/websockify')
+WEBSOCKIFY_TARGET_FILE = environ.get('PROXSTAR_WEBSOCKIFY_TARGET_FILE', '/opt/app-root/src/targets')
diff --git a/proxstar/__init__.py b/proxstar/__init__.py
index ffb3953..2e10090 100644
--- a/proxstar/__init__.py
+++ b/proxstar/__init__.py
@@ -1,5 +1,6 @@
 import os
 import time
+import atexit
 import subprocess
 from rq import Queue
 from redis import Redis
@@ -10,6 +11,7 @@ from werkzeug.contrib.cache import SimpleCache
 from flask_pyoidc.flask_pyoidc import OIDCAuthentication
 from flask import Flask, render_template, request, redirect, send_from_directory, session
 from proxstar.db import *
+from proxstar.vnc import *
 from proxstar.util import *
 from proxstar.tasks import *
 from proxstar.starrs import *
@@ -28,6 +30,13 @@ app.config.from_pyfile(config)
 app.config["GIT_REVISION"] = subprocess.check_output(
     ['git', 'rev-parse', '--short', 'HEAD']).decode('utf-8').rstrip()
 
+with open('proxmox_ssh_key', 'w') as key:
+    key.write(app.config['PROXMOX_SSH_KEY'])
+
+start_websockify(app.config['WEBSOCKIFY_PATH'], app.config['WEBSOCKIFY_TARGET_FILE'])
+
+ssh_tunnels = []
+
 retry = 0
 while retry < 5:
     try:
@@ -132,6 +141,7 @@ def vm_details(vmid):
         vm = get_vm(proxmox, vmid)
         vm['vmid'] = vmid
         vm['config'] = get_vm_config(proxmox, vmid)
+        vm['node'] = get_vm_node(proxmox, vmid)
         vm['disks'] = get_vm_disks(proxmox, vmid, config=vm['config'])
         vm['iso'] = get_vm_iso(proxmox, vmid, config=vm['config'])
         vm['interfaces'] = []
@@ -179,6 +189,23 @@ def vm_power(vmid, action):
         return '', 403
 
 
+@app.route("/vm/<string:vmid>/console", methods=['POST'])
+@auth.oidc_auth
+def vm_console(vmid):
+    user = session['userinfo']['preferred_username']
+    rtp = 'rtp' in session['userinfo']['groups']
+    proxmox = connect_proxmox()
+    if rtp or int(vmid) in get_user_allowed_vms(proxmox, db, user):
+        port = str(5900 + int(vmid))
+        token = add_vnc_target(port)
+        node = "{}.csh.rit.edu".format(get_vm_node(proxmox, vmid))
+        tunnel = start_ssh_tunnel(node, port)
+        ssh_tunnels.append(tunnel)
+        return token, 200
+    else:
+        return '', 403
+
+
 @app.route("/vm/<string:vmid>/cpu/<int:cores>", methods=['POST'])
 @auth.oidc_auth
 def vm_cpu(vmid, cores):
@@ -420,23 +447,23 @@ def template_disk(template_id):
     return get_template_disk(db, template_id)
 
 
-@app.route('/vm/<string:vmid>/rrd/<path:path>')
-@auth.oidc_auth
-def send_rrd(vmid, path):
-    return send_from_directory("rrd/{}".format(vmid), path)
-
-
-@app.route('/novnc/<path:path>')
-@auth.oidc_auth
-def send_novnc(path):
-    return send_from_directory('static/novnc-pve/novnc', path)
-
-
 @app.route("/logout")
 @auth.oidc_logout
 def logout():
     return redirect(url_for('list_vms'), 302)
 
 
+def exit_handler():
+    stop_websockify()
+    for tunnel in ssh_tunnels:
+        try:
+            tunnel.stop()
+        except:
+            pass
+
+
+atexit.register(exit_handler)
+
+
 if __name__ == "__main__":
     app.run()
diff --git a/proxstar/proxmox.py b/proxstar/proxmox.py
index fd9eb52..e993dc6 100644
--- a/proxstar/proxmox.py
+++ b/proxstar/proxmox.py
@@ -287,12 +287,6 @@ def get_pools(proxmox, db):
     return pools
 
 
-def get_rrd_for_vm(proxmox, vmid, source, time):
-    node = proxmox.nodes(get_vm_node(proxmox, vmid))
-    image = node.qemu(vmid).rrd.get(ds=source, timeframe=time)['image']
-    return image
-
-
 def delete_user_pool(proxmox, pool):
     proxmox.pools(pool).delete()
     users = proxmox.access.users.get()
diff --git a/proxstar/static/css/styles.css b/proxstar/static/css/styles.css
index b662c38..b637e03 100644
--- a/proxstar/static/css/styles.css
+++ b/proxstar/static/css/styles.css
@@ -105,13 +105,6 @@ table, th, td {
     margin: 0px auto;
 }
 
-.rrd-graph {
-    width: 800px;
-    height: 200px;
-    padding-top: 10px;
-    padding-bottom: 10px;
-}
-
 .user-img {
     border-radius: 50%;
     height: 25px;
diff --git a/proxstar/static/js/script.js b/proxstar/static/js/script.js
index fe599e7..eed7f4f 100644
--- a/proxstar/static/js/script.js
+++ b/proxstar/static/js/script.js
@@ -802,3 +802,23 @@ function hide_for_template(obj) {
         hide_area.style.display = 'none';
     }
 }
+
+$("#console-vm").click(function(){
+    const vmname = $(this).data('vmname');
+    const vmid = $(this).data('vmid');
+    fetch(`/vm/${vmid}/console`, {
+        credentials: 'same-origin',
+        method: 'post'
+    }).then((response) => {
+        return response.text()
+    }).then((token) => {
+        window.location = `/static/noVNC/vnc.html?autoconnect=true&?encrypt=true&?host=proxstar.csh.rit.edu&?port=8081&?path=path?token=${token}`;
+    }).catch(err => {
+        if (err) {
+            swal("Uh oh...", `Unable to start console for ${vmname}. Please try again later.`, "error");
+        } else {
+            swal.stopLoading();
+            swal.close();
+        }
+    });
+});
diff --git a/proxstar/static/noVNC b/proxstar/static/noVNC
new file mode 160000
index 0000000..37b4d13
--- /dev/null
+++ b/proxstar/static/noVNC
@@ -0,0 +1 @@
+Subproject commit 37b4d13db81e0e80e117c07b86ff98714c7b6b1a
diff --git a/proxstar/templates/vm_details.html b/proxstar/templates/vm_details.html
index 2627dba..84a59d1 100644
--- a/proxstar/templates/vm_details.html
+++ b/proxstar/templates/vm_details.html
@@ -21,7 +21,7 @@
                     {% endif %}
                     {% if vm['qmpstatus'] == 'running' or vm['qmpstatus'] == 'paused' %}
                     {% if vm['qmpstatus'] == 'running' %}
-                    <a href="https://proxmox01.csh.rit.edu/#v1:0:=qemu%2F{{ vm['vmid'] }}:4::::8::" target="_blank">Open VM Console</a>
+                    <button class="btn btn-success proxstar-actionbtn" id="console-vm" name="console" data-vmid="{{ vm['vmid'] }}" data-vmname="{{ vm['name'] }}">CONSOLE</button>
                     <button class="btn btn-info proxstar-actionbtn" id="suspend-vm" name="suspend" data-vmid="{{ vm['vmid'] }}" data-vmname="{{ vm['name'] }}">SUSPEND</button>
                     {% endif %}
                     <button class="btn btn-info proxstar-actionbtn" id="shutdown-vm" name="shutdown" data-vmid="{{ vm['vmid'] }}" data-vmname="{{ vm['name'] }}">SHUTDOWN</button>
@@ -82,6 +82,8 @@
                         <dd>{{ vm['vmid'] }}</dd>
                         <dt>Status</dt>
                         <dd>{{ vm['qmpstatus'] }}</dd>
+                        <dt>Node</dt>
+                        <dd>{{ vm['node'] }}</dd>
                         <dt>Cores</dt>
                         <dd>
                             {{ vm['config']['cores'] * vm['config'].get('sockets', 1) }}
diff --git a/proxstar/vnc.py b/proxstar/vnc.py
new file mode 100644
index 0000000..378c377
--- /dev/null
+++ b/proxstar/vnc.py
@@ -0,0 +1,71 @@
+import time
+import subprocess
+from sshtunnel import SSHTunnelForwarder
+from proxstar.util import *
+from flask import current_app as app
+
+
+def start_websockify(websockify_path, target_file):
+    result = subprocess.run(['pgrep', 'websockify'], stdout=subprocess.PIPE)
+    if not result.stdout:
+        subprocess.call(
+            [
+                websockify_path, '8081', '--token-plugin',
+                'TokenFile', '--token-source', target_file,
+                '-D'
+            ],
+            stdout=subprocess.PIPE)
+
+
+def stop_websockify():
+    result = subprocess.run(['pgrep', 'websockify'], stdout=subprocess.PIPE)
+    if result.stdout:
+        pid = result.stdout.strip()
+        subprocess.run(['kill', pid], stdout=subprocess.PIPE)
+        time.sleep(3)
+        if subprocess.run(
+            ['pgrep', 'websockify'], stdout=subprocess.PIPE).stdout:
+            time.sleep(10)
+            if subprocess.run(
+                ['pgrep', 'websockify'], stdout=subprocess.PIPE).stdout:
+                print('Websockify didn\'t stop, killing forcefully.')
+                subprocess.run(['kill', '-9', pid], stdout=subprocess.PIPE)
+
+
+def get_vnc_targets():
+    target_file = open(app.config['WEBSOCKIFY_TARGET_FILE'])
+    targets = []
+    for line in target_file:
+        target_dict = dict()
+        values = line.strip().split(':')
+        target_dict['token'] = values[0]
+        target_dict['port'] = values[2]
+        targets.append(target_dict)
+    target_file.close()
+    return (targets)
+
+
+def add_vnc_target(port):
+    targets = get_vnc_targets()
+    target = next((target for target in targets if target['port'] == port),
+                  None)
+    if target:
+        return target['token']
+    else:
+        target_file = open(app.config['WEBSOCKIFY_TARGET_FILE'], 'a')
+        token = gen_password(32, 'abcdefghijklmnopqrstuvwxyz0123456789')
+        target_file.write("{}: 127.0.0.1:{}\n".format(token, str(port)))
+        target_file.close()
+        return token
+
+
+def start_ssh_tunnel(node, port):
+    port = int(port)
+    server = SSHTunnelForwarder(
+        node,
+        ssh_username='root',
+        ssh_pkey='proxmox_ssh_key',
+        remote_bind_address=('127.0.0.1', port),
+        local_bind_address=('127.0.0.1', port))
+    server.start()
+    return server
diff --git a/requirements.txt b/requirements.txt
index f9d2e47..eb29e34 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -10,3 +10,4 @@ rq-scheduler==0.7.0
 gunicorn
 raven
 paramiko
+websockify