From 7c2b9df0d03207c7a884754ed2fede0984aa0cb2 Mon Sep 17 00:00:00 2001
From: viveksantayana <me@viveksantayana.co.uk>
Date: Sat, 4 Dec 2021 17:40:01 +0000
Subject: [PATCH] Added custom 404 display and login redirect

---
 ref-test/admin/models/users.py                    | 10 ++++++++--
 ref-test/admin/static/js/script.js                |  7 ++++++-
 ref-test/admin/views.py                           |  3 ++-
 ref-test/main.py                                  |  6 +++++-
 ref-test/quiz/templates/quiz/404.html             |  8 ++++++++
 ref-test/quiz/templates/quiz/components/base.html |  8 ++++----
 6 files changed, 33 insertions(+), 9 deletions(-)
 create mode 100644 ref-test/quiz/templates/quiz/404.html

diff --git a/ref-test/admin/models/users.py b/ref-test/admin/models/users.py
index 5db7156..ea7824e 100644
--- a/ref-test/admin/models/users.py
+++ b/ref-test/admin/models/users.py
@@ -1,4 +1,4 @@
-from flask import flash, make_response, Response
+from flask import flash, make_response, Response, session
 from flask.helpers import url_for
 from flask.json import jsonify
 from werkzeug.security import generate_password_hash, check_password_hash
@@ -61,7 +61,13 @@ class User:
             return jsonify({ 'error': f'Username {self.username} does not exist.' }), 401
         if not check_password_hash( user['password'], self.password ):
             return jsonify({ 'error': f'The password you entered is incorrect.' }), 401
-        resp = make_response(jsonify({ 'success': f'Successfully logged in user {self.username}.' }), 200)
+        response = {
+            'success': f'Successfully logged in user {self.username}.'
+        }
+        if 'prev_page' in session:
+            response['redirect_to'] = session['prev_page']
+            session.pop('prev_page')
+        resp = make_response(jsonify(response), 200)
         self._id = user['_id']
         self.start_session(resp)
         return resp
diff --git a/ref-test/admin/static/js/script.js b/ref-test/admin/static/js/script.js
index 00c3b45..d950bc3 100644
--- a/ref-test/admin/static/js/script.js
+++ b/ref-test/admin/static/js/script.js
@@ -27,7 +27,12 @@ $('form.form-post').submit(function(event) {
         data: data,
         dataType: 'json',
         success: function(response) {
-            window.location.href = rel_success;
+            if (response.redirect_to) {
+                window.location.href = response.redirect_to;
+            }
+            else {
+                window.location.href = rel_success;
+            }
         },
         error: function(response) {
             error_response(response);
diff --git a/ref-test/admin/views.py b/ref-test/admin/views.py
index 1ab078b..6bc2523 100644
--- a/ref-test/admin/views.py
+++ b/ref-test/admin/views.py
@@ -1,4 +1,4 @@
-from flask import Blueprint, render_template, flash, redirect, request, jsonify, abort, make_response
+from flask import Blueprint, render_template, flash, redirect, request, jsonify, abort, session
 from flask.helpers import url_for
 from functools import wraps
 from datetime import datetime
@@ -56,6 +56,7 @@ def login_required(function):
     @wraps(function)
     def decorated_function(*args, **kwargs):
         if not check_login():
+            session['prev_page'] = request.url
             flash('Please log in to view this page.', 'alert')
             return redirect(url_for('admin_auth.login'))
         return function(*args, **kwargs)
diff --git a/ref-test/main.py b/ref-test/main.py
index 47b77d1..bf1899a 100644
--- a/ref-test/main.py
+++ b/ref-test/main.py
@@ -1,6 +1,6 @@
 from datetime import datetime
 
-from flask import Flask, flash, request
+from flask import Flask, flash, request, render_template
 from flask.helpers import url_for
 from flask.json import jsonify
 from flask_bootstrap import Bootstrap
@@ -77,4 +77,8 @@ if __name__ == '__main__':
     def _get_id_from_cookie():
         return dict(get_id_from_cookie = get_id_from_cookie)
 
+    @app.errorhandler(404)
+    def _404_handler(e):
+        return render_template('/quiz/404.html'), 404
+
     app.run(host=app.config['APP_HOST'])
\ No newline at end of file
diff --git a/ref-test/quiz/templates/quiz/404.html b/ref-test/quiz/templates/quiz/404.html
new file mode 100644
index 0000000..2601289
--- /dev/null
+++ b/ref-test/quiz/templates/quiz/404.html
@@ -0,0 +1,8 @@
+{% extends "quiz/components/base.html" %}
+
+{% block content %}
+    <h1>Page Not Found</h1>
+    <p>
+        The page you were looking for does not exist. Try going back and navigating to the desired destination correctly.
+    </p>
+{% endblock %}
\ No newline at end of file
diff --git a/ref-test/quiz/templates/quiz/components/base.html b/ref-test/quiz/templates/quiz/components/base.html
index fbc0ed6..ccd4f07 100644
--- a/ref-test/quiz/templates/quiz/components/base.html
+++ b/ref-test/quiz/templates/quiz/components/base.html
@@ -30,11 +30,11 @@
                 {% include "quiz/components/server-alerts.html" %}
             {% endblock %}
             {% block content %}{% endblock %}
-        </div>
 
-        <footer class="container site-footer">
-            {% include "quiz/components/footer.html" %}
-        </footer>
+            <footer class="container site-footer">
+                {% include "quiz/components/footer.html" %}
+            </footer>
+        </div>
 
         <!-- JQuery, Popper, and Bootstrap js dependencies -->
         <script