From 635e4c0f9f0a80c3dd7744f899200cac89ddd37d Mon Sep 17 00:00:00 2001
From: Jordan Rodgers <com6056@gmail.com>
Date: Tue, 6 Mar 2018 20:49:49 -0500
Subject: [PATCH] add rq dashboard for RTPs

---
 proxstar/__init__.py         | 21 +++++++++++++++++++--
 proxstar/templates/base.html |  6 ++++++
 proxstar/user.py             |  3 ++-
 requirements.txt             |  1 +
 4 files changed, 28 insertions(+), 3 deletions(-)

diff --git a/proxstar/__init__.py b/proxstar/__init__.py
index 08fbe62..d4f5625 100644
--- a/proxstar/__init__.py
+++ b/proxstar/__init__.py
@@ -4,6 +4,7 @@ import psutil
 import atexit
 import psycopg2
 import subprocess
+import rq_dashboard
 from rq import Queue
 from redis import Redis
 from rq_scheduler import Scheduler
@@ -20,6 +21,7 @@ from proxstar.ldapdb import *
 from proxstar.proxmox import *
 
 app = Flask(__name__)
+app.config.from_object(rq_dashboard.default_settings)
 if os.path.exists(
         os.path.join(
             app.config.get('ROOT_DIR', os.getcwd()), "config.local.py")):
@@ -84,6 +86,19 @@ if 'cleanup_vnc' not in scheduler:
         interval=3600)
 
 
+def add_rq_dashboard_auth(blueprint):
+    @blueprint.before_request
+    @auth.oidc_auth
+    def rq_dashboard_auth(*args, **kwargs):
+        if 'rtp' not in session['userinfo']['groups']:
+            return '', 403
+
+
+rq_dashboard_blueprint = rq_dashboard.blueprint
+add_rq_dashboard_auth(rq_dashboard_blueprint)
+app.register_blueprint(rq_dashboard_blueprint, url_prefix="/rq")
+
+
 @app.route("/")
 @app.route("/user/<string:user_view>")
 @auth.oidc_auth
@@ -97,7 +112,8 @@ def list_vms(user_view=None):
         user_view = User(user_view)
         vms = user_view.vms
         for pending_vm in user_view.pending_vms:
-            vm = next((vm for vm in vms if vm['name'] == pending_vm['name']), None)
+            vm = next((vm for vm in vms if vm['name'] == pending_vm['name']),
+                      None)
             if vm:
                 vms[vms.index(vm)]['status'] = pending_vm['status']
                 vms[vms.index(vm)]['pending'] = True
@@ -111,7 +127,8 @@ def list_vms(user_view=None):
         if user.active:
             vms = user.vms
             for pending_vm in user.pending_vms:
-                vm = next((vm for vm in vms if vm['name'] == pending_vm['name']), None)
+                vm = next((vm for vm in vms
+                           if vm['name'] == pending_vm['name']), None)
                 if vm:
                     vms[vms.index(vm)]['status'] = pending_vm['status']
                     vms[vms.index(vm)]['pending'] = True
diff --git a/proxstar/templates/base.html b/proxstar/templates/base.html
index d1acd67..ffc1436 100644
--- a/proxstar/templates/base.html
+++ b/proxstar/templates/base.html
@@ -41,6 +41,12 @@
                         </a>
                     </li>
                     {% if user.rtp %}
+                    <li>
+                        <a href="/rq">
+                            <span class="glyphicon glyphicon-eye-open"></span>
+                            RQ Dashboard
+                        </a>
+                    </li>
                     <li>
                         <a href="/settings">
                             <span class="glyphicon glyphicon-cog"></span>
diff --git a/proxstar/user.py b/proxstar/user.py
index 5aa1904..dfa19b9 100644
--- a/proxstar/user.py
+++ b/proxstar/user.py
@@ -33,7 +33,8 @@ class User(object):
 
     @lazy_property
     def pending_vms(self):
-        jobs = StartedJobRegistry('default', connection=redis_conn).get_job_ids()
+        jobs = StartedJobRegistry(
+            'default', connection=redis_conn).get_job_ids()
         for job_id in q.job_ids:
             jobs.append(job_id)
         pending_vms = []
diff --git a/requirements.txt b/requirements.txt
index 7738f82..f6096f9 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -14,3 +14,4 @@ websockify
 sshtunnel
 psutil
 requests
+rq_dashboard