From df100a25b8b164b8d533830f4d7423d38ab24bf8 Mon Sep 17 00:00:00 2001
From: Jordan Rodgers <com6056@gmail.com>
Date: Wed, 7 Mar 2018 01:49:42 -0500
Subject: [PATCH] add proper 404 and 403 pages

---
 proxstar/__init__.py        | 22 +++++++++++++++++-----
 proxstar/templates/403.html |  8 ++++++++
 proxstar/templates/404.html |  8 ++++++++
 3 files changed, 33 insertions(+), 5 deletions(-)
 create mode 100644 proxstar/templates/403.html
 create mode 100644 proxstar/templates/404.html

diff --git a/proxstar/__init__.py b/proxstar/__init__.py
index 53d4834..da5925a 100644
--- a/proxstar/__init__.py
+++ b/proxstar/__init__.py
@@ -11,7 +11,7 @@ from rq_scheduler import Scheduler
 from sqlalchemy import create_engine
 from sqlalchemy.orm import sessionmaker
 from flask_pyoidc.flask_pyoidc import OIDCAuthentication
-from flask import Flask, render_template, request, redirect, session
+from flask import Flask, render_template, request, redirect, session, abort
 from proxstar.db import *
 from proxstar.vm import VM
 from proxstar.vnc import *
@@ -91,7 +91,7 @@ def add_rq_dashboard_auth(blueprint):
     @auth.oidc_auth
     def rq_dashboard_auth(*args, **kwargs):
         if 'rtp' not in session['userinfo']['groups']:
-            return '', 403
+            abort(403)
 
 
 rq_dashboard_blueprint = rq_dashboard.blueprint
@@ -99,6 +99,18 @@ add_rq_dashboard_auth(rq_dashboard_blueprint)
 app.register_blueprint(rq_dashboard_blueprint, url_prefix="/rq")
 
 
+@app.errorhandler(404)
+def not_found(e):
+    user = User(session['userinfo']['preferred_username'])
+    return render_template('404.html', user=user), 404
+
+
+@app.errorhandler(403)
+def forbidden(e):
+    user = User(session['userinfo']['preferred_username'])
+    return render_template('403.html', user=user), 403
+
+
 @app.route("/")
 @app.route("/user/<string:user_view>")
 @auth.oidc_auth
@@ -107,7 +119,7 @@ def list_vms(user_view=None):
     rtp_view = False
     proxmox = connect_proxmox()
     if user_view and not user.rtp:
-        return '', 403
+        abort(403)
     elif user_view and user.rtp:
         user_view = User(user_view)
         vms = user_view.vms
@@ -190,7 +202,7 @@ def vm_details(vmid):
             limits=user.limits,
             usage_check=usage_check)
     else:
-        return '', 403
+        abort(403)
 
 
 @app.route("/vm/<string:vmid>/power/<string:action>", methods=['POST'])
@@ -479,7 +491,7 @@ def settings():
             ignored_pools=ignored_pools,
             allowed_users=allowed_users)
     else:
-        return '', 403
+        abort(403)
 
 
 @app.route("/pool/<string:pool>/ignore", methods=['POST', 'DELETE'])
diff --git a/proxstar/templates/403.html b/proxstar/templates/403.html
new file mode 100644
index 0000000..ae545cd
--- /dev/null
+++ b/proxstar/templates/403.html
@@ -0,0 +1,8 @@
+{% extends "base.html" %}
+{% block body %}
+
+<div class="container text-center">
+    <h1>You aren't allowed to view this page!</h1>
+</div>
+
+{% endblock %}
diff --git a/proxstar/templates/404.html b/proxstar/templates/404.html
new file mode 100644
index 0000000..43f0265
--- /dev/null
+++ b/proxstar/templates/404.html
@@ -0,0 +1,8 @@
+{% extends "base.html" %}
+{% block body %}
+
+<div class="container text-center">
+    <h1>That page doesn't exist!</h1>
+</div>
+
+{% endblock %}